# --------------------------------------------------------------
# app.py
# --------------------------------------------------------------
# 1️⃣  Load the Qwen2.5‑1.5B‑Instruct model (float16, 4‑bit quantised)
# 2️⃣  Expose two things:
#      • OpenAI‑compatible endpoint   → /v1/chat/completions
#      • Static UI (HTML/CSS/JS)      → /
# --------------------------------------------------------------

import os
import torch
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse, HTMLResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field
from transformers import AutoModelForCausalLM, AutoTokenizer
import uvicorn
import time

# --------------------------------------------------------------
# 1️⃣  Model loading (once, shared by all requests)
# --------------------------------------------------------------
MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"

print("🚀 Loading model – this may take a minute …")
tokenizer = AutoTokenizer.from_pretrained(
    MODEL_NAME,
    trust_remote_code=True,
    use_fast=True
)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16,          # float16 works on the free tier (CPU‑only)
    device_map="auto",
    trust_remote_code=True
)

# --------------------------------------------------------------
# 2️⃣  FastAPI app
# --------------------------------------------------------------
app = FastAPI(
    title="Qwen2.5‑1.5B‑Instruct (Traditional Chinese) – OpenAI‑compatible",
    description="Chat endpoint (`/v1/chat/completions`) + modern web UI (`/`).",
    version="0.1.0"
)

# Serve static files from ./static
app.mount("/static", StaticFiles(directory="static"), name="static")

# --------------------------------------------------------------
# 3️⃣  Pydantic models – OpenAI‑compatible request/response
# --------------------------------------------------------------
class Message(BaseModel):
    role: str = Field(..., description="`system`, `user` or `assistant`")
    content: str

class ChatRequest(BaseModel):
    model: str = Field(default="qwen2.5-1.5b-instruct")
    messages: list[Message]
    temperature: float = Field(default=0.7, ge=0.0, le=2.0)
    max_tokens: int = Field(default=512, ge=1, le=2048)
    top_p: float = Field(default=0.9, ge=0.0, le=1.0)
    stream: bool = Field(default=False)

class Choice(BaseModel):
    index: int = 0
    message: Message
    finish_reason: str = "stop"

class Usage(BaseModel):
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int

class ChatResponse(BaseModel):
    id: str = "chatcmpl-123"
    object: str = "chat.completion"
    created: int = Field(default_factory=lambda: int(time.time()))
    model: str
    choices: list[Choice]
    usage: Usage

# --------------------------------------------------------------
# 4️⃣  Helper – generate a response from the model
# --------------------------------------------------------------
def _generate(request: ChatRequest) -> str:
    # Convert list[Message] → the format the tokenizer expects
    msgs = [{"role": m.role, "content": m.content} for m in request.messages]

    # Build the prompt using the model’s chat template
    prompt = tokenizer.apply_chat_template(
        msgs,
        tokenize=False,
        add_generation_prompt=True
    )

    # Tokenise
    inputs = tokenizer([prompt], return_tensors="pt").to(model.device)

    # Generation
    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=request.max_tokens,
            temperature=request.temperature,
            top_p=request.top_p,
            do_sample=True,
            repetition_penalty=1.1,
            pad_token_id=tokenizer.eos_token_id,
        )

    # Remove the prompt part
    generated = output_ids[0][len(inputs.input_ids[0]):]
    text = tokenizer.decode(generated, skip_special_tokens=True)
    return text

# --------------------------------------------------------------
# 5️⃣  OpenAI‑compatible endpoint
# --------------------------------------------------------------
@app.post("/v1/chat/completions", response_model=ChatResponse)
async def chat_completions(req: ChatRequest):
    try:
        answer = _generate(req)

        # Token counts (approximate)
        prompt_tokens = len(tokenizer.encode(req.messages[-1].content))
        completion_tokens = len(tokenizer.encode(answer))
        total_tokens = prompt_tokens + completion_tokens

        response = ChatResponse(
            model=req.model,
            choices=[
                Choice(
                    index=0,
                    message=Message(role="assistant", content=answer),
                    finish_reason="stop"
                )
            ],
            usage=Usage(
                prompt_tokens=prompt_tokens,
                completion_tokens=completion_tokens,
                total_tokens=total_tokens,
            ),
        )
        return response
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc))

# --------------------------------------------------------------
# 6️⃣  UI entry point – serve the HTML page
# --------------------------------------------------------------
@app.get("/", response_class=HTMLResponse)
async def index():
    # The file lives in ./static/index.html
    with open(os.path.join("static", "index.html"), encoding="utf-8") as f:
        return HTMLResponse(f.read())

# --------------------------------------------------------------
# 7️⃣  Run locally (useful for debugging)
# --------------------------------------------------------------
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)