from fastapi import FastAPI, Depends, HTTPException, Query
from transformers import AutoModelForCausalLM, AutoTokenizer
from typing import List
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles

app = FastAPI()

app.mount("/static", StaticFiles(directory="static"), name="static")

@app.get("/", response_class=HTMLResponse)
async def read_root():
    with open("static/index.html", "r") as f:
        content = f.read()
    return HTMLResponse(content=content)

# Initialize model and tokenizer
# tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen-14B-Chat-int4")
# model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen-14B-Chat-int4").eval()

@app.post("/chat/")
def chat(user_input: str, api_key: str):
    # Here you can validate the API key, e.g., check if it exists in your database
    # If the API key is not valid, raise an HTTPException
    # if not validate_api_key(api_key):
    #     raise HTTPException(status_code=400, detail="Invalid API key")

    # Tokenize the user input and get model's response
    # input_ids = tokenizer.encode(user_input, return_tensors="pt")
    # output = model.generate(input_ids)
    # response = tokenizer.decode(output[0], skip_special_tokens=True)
    
    return {"response": f"user input: {user_input}, api_key: {api_key}"}