from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

app = FastAPI()

# Load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")

# In-memory history per user
chat_history = {}

@app.get("/")
async def root():
    return {"message": "🟢 API is running. Use /ai?query=Hello&user_id=yourname"}

@app.get("/ai")
async def chat(request: Request):
    query_params = dict(request.query_params)
    user_input = query_params.get("query", "")
    user_id = query_params.get("user_id", "default")

    if not user_input:
        return JSONResponse({"error": "Missing 'query' parameter"}, status_code=400)

    new_input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors='pt')
    user_history = chat_history.get(user_id, [])
    bot_input_ids = torch.cat(user_history + [new_input_ids], dim=-1) if user_history else new_input_ids

    output_ids = model.generate(bot_input_ids, max_new_tokens=100, pad_token_id=tokenizer.eos_token_id)
    response = tokenizer.decode(output_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)

    chat_history[user_id] = [bot_input_ids, output_ids]
    return JSONResponse({"reply": response})

# Only needed if running locally, not in Hugging Face Space
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)