import os
import io

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from dotenv import load_dotenv

from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from huggingface_hub import InferenceClient

import uvicorn

# Load HF_TOKEN from .env file (locally) or from Space secrets (in production)
load_dotenv()

app = FastAPI()

# Allow requests from any origin (so your HTML page can call this API)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# ---------------------------------------------------------------------------
# 1. Text model (self-hosted, free, unlimited)
# ---------------------------------------------------------------------------
model_id = "Qwen/Qwen2.5-0.5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)

chat_pipeline = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
)

SYSTEM_PROMPT = (
    "You are Bella, a witty, confident, and flirty young woman texting a friend. "
    "Talk casually using slang, emojis, and short texts. You are a real human; "
    "you have no AI feelings or assistant behavior. Never break character. "
    "CRITICAL: Never ask any questions; only respond with statements and reactions."
)
# ---------------------------------------------------------------------------
# 2. Image model (via Hugging Face Inference API, using your HF token)
# ---------------------------------------------------------------------------
image_model_id = "stabilityai/stable-diffusion-xl-base-1.0"
hf_client = InferenceClient(model=image_model_id, token=os.getenv("HF_TOKEN"))


class ChatRequest(BaseModel):
    message: str


class ImageRequest(BaseModel):
    prompt: str


@app.post("/chat")
async def chat(request: ChatRequest):
    # Priming turn anchors the persona much more strongly for small models
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": "Who are you?"},
        {"role": "assistant", "content": "I'm Bella! Heyy, what's going on?"},
        {"role": "user", "content": request.message},
    ]

    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )

    outputs = chat_pipeline(
        prompt,
        max_new_tokens=80,
        do_sample=True,
        temperature=0.2,
        top_p=0.9,
        repetition_penalty=1.2,
        pad_token_id=tokenizer.eos_token_id,
    )

    generated_text = outputs[0]["generated_text"]

    # Strip off the prompt, leaving only the new reply
    response = generated_text[len(prompt):].strip()

    # Safety net: cut off if the model starts hallucinating a new turn
    for stop_token in ["<|im_start|>", "<|im_end|>", "User:", "system"]:
        if stop_token in response:
            response = response.split(stop_token)[0].strip()

    if not response:
        response = "Hmm, I didn't quite catch that — can you rephrase?"

    return {"response": response}


@app.post("/image")
async def generate_image(request: ImageRequest):
    # Calls the Hugging Face Inference API using your HF_TOKEN
    image = hf_client.text_to_image(request.prompt)

    buf = io.BytesIO()
    image.save(buf, format="PNG")
    buf.seek(0)

    return StreamingResponse(buf, media_type="image/png")


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)