from fastapi import FastAPI, Request, Form
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.templating import Jinja2Templates
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

app = FastAPI()

# Configura las plantillas Jinja2
templates = Jinja2Templates(directory="templates")

# Define el personaje
personaje = "rias"
user="user"
chat={
    personaje: f"hola soy {personaje} no esperaba verte por aqui",
    user:f"hola "
}
# Monta la carpeta 'static' para servir archivos estáticos
app.mount("/static", StaticFiles(directory="static"), name="static")

# Ruta para mostrar los personajes
@app.get("/", response_class=HTMLResponse)
async def read_html(request: Request):
    return templates.TemplateResponse("listapersonajes.html", {"request": request})

# Ruta dinámica para cada personaje
@app.get("/personajes/{personaje}", response_class=HTMLResponse)
async def personaje_detalle(request: Request, personaje: str):
    # El contexto es el nombre de la imagen que se usará
    context = {
        "character_image": f"{personaje}.jpg" , # Asume que el nombre de la imagen es igual al personaje
        "character_name": personaje.capitalize()  # Nombre del personaje con la primera letra en mayúscula

    }
    return templates.TemplateResponse("chat.html", {"request": request, **context})






model_name = "nakodanei/Blue-Orchid-2x7b"

# Inicialización global
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    torch_dtype=torch.float16
)

@app.post("/personajes/{personaje}/chat", response_class=HTMLResponse)
async def chat_with_character(request: Request, personaje: str, user_input: str):
    prompt = f"""
    ### Instruction:
    You are {personaje}, a sexy girl who has been dating the user for 2 months.

    ### Input:
    User: {user_input}

    ### Response:
    Character:
    """
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
    outputs = model.generate(
        **inputs,
        max_new_tokens=500,
        pad_token_id=tokenizer.eos_token_id,
        do_sample=True,
        temperature=0.7
    )
    generated_response = tokenizer.decode(outputs[0], skip_special_tokens=True)

    return {"response": generated_response}