from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    TextIteratorStreamer
)
import torch
import uvicorn
import threading
import json

# =========================
# APP
# =========================

app = FastAPI()

stop_flags = {}

# =========================
# MODEL
# =========================

MODEL_ID = "Qwen/Qwen2.5-Coder-1.5B-Instruct"

print("🚀 Loading Fast Coder Model...")

device = torch.device(
    "cuda" if torch.cuda.is_available() else "cpu"
)

torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

# =========================
# TOKENIZER
# =========================

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_ID,
    trust_remote_code=True
)

# =========================
# MODEL
# =========================

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    trust_remote_code=True,
    torch_dtype=torch.float16 if device.type == "cuda" else torch.float32
)

model = model.to(device)
model.eval()

print(f"✅ Loaded on {device}")

# =========================
# REQUEST
# =========================

class ChatRequest(BaseModel):
    message: str
    conversation_id: str
    temperature: float = 0.1

# =========================
# SYSTEM PROMPT
# =========================

SYSTEM_PROMPT = """
You are a strict expert programming assistant.

CRITICAL RULES:
- Answer ONLY the user's latest request
- NEVER continue conversations
- NEVER generate extra examples unless asked
- NEVER explain unnecessarily
- NEVER repeat code
- NEVER simulate dialogue
- ALWAYS close markdown code blocks properly
- ALWAYS return complete executable code
- Stop immediately after final answer

CODE RULES:
- Use proper markdown
- Use ```language
- Keep formatting clean
- No duplicate code
- No unfinished code
"""

# =========================
# STOP WORDS
# =========================

STOP_WORDS = [
    "<|im_end|>",
    "<|endoftext|>",
    "<|eot_id|>",
    "User:",
    "Assistant:",
    "Human:"
]

# =========================
# CLEAN OUTPUT
# =========================

def clean_output(text):

    for w in STOP_WORDS:

        if w in text:
            text = text.split(w)[0]

    return text.strip()

# =========================
# BUILD INPUTS
# =========================

def build_inputs(message):

    messages = [
        {
            "role": "system",
            "content": SYSTEM_PROMPT
        },
        {
            "role": "user",
            "content": message
        }
    ]

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

    return tokenizer(
        text,
        return_tensors="pt"
    ).to(device)

# =========================
# STOP ENDPOINT
# =========================

@app.post("/v1/stop")
def stop(data: dict):

    stop_flags[data.get("conversation_id")] = True

    return {
        "status": "stopped"
    }

# =========================
# NORMAL CHAT
# =========================

@app.post("/v1/chat")
def chat(req: ChatRequest):

    inputs = build_inputs(req.message)

    with torch.inference_mode():

        output = model.generate(
            **inputs,
            max_new_tokens=512,
            do_sample=False,
            temperature=req.temperature,
            top_p=1.0,
            repetition_penalty=1.08,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id
        )

    result = tokenizer.decode(
        output[0][inputs.input_ids.shape[1]:],
        skip_special_tokens=True
    )

    result = clean_output(result)

    return {
        "response": result
    }

# =========================
# STREAM CHAT
# =========================

@app.post("/v1/chat/stream")
def stream_chat(req: ChatRequest):

    inputs = build_inputs(req.message)

    streamer = TextIteratorStreamer(
        tokenizer,
        skip_prompt=True,
        skip_special_tokens=True
    )

    generation_kwargs = dict(
        **inputs,
        streamer=streamer,
        max_new_tokens=512,
        do_sample=False,
        temperature=req.temperature,
        top_p=1.0,
        repetition_penalty=1.08,
        pad_token_id=tokenizer.eos_token_id,
        eos_token_id=tokenizer.eos_token_id
    )

    thread = threading.Thread(
        target=model.generate,
        kwargs=generation_kwargs
    )

    thread.start()

    def generate():

        full_text = ""

        for token in streamer:

            if stop_flags.get(req.conversation_id):

                stop_flags[req.conversation_id] = False
                break

            if not token:
                continue

            stop_hit = False

            for sw in STOP_WORDS:

                if sw in token:
                    token = token.split(sw)[0]
                    stop_hit = True
                    break

            if token:

                full_text += token

                # stop after completed markdown block
                if full_text.count("```") >= 2:
                    yield f"data: {json.dumps({'choices':[{'delta':{'content': token}}]})}\n\n"
                    break

                yield f"data: {json.dumps({'choices':[{'delta':{'content': token}}]})}\n\n"

            if stop_hit:
                break

        full_text = clean_output(full_text)

        yield "event: done\ndata: {}\n\n"
        yield "data: [DONE]\n\n"

    return StreamingResponse(
        generate(),
        media_type="text/event-stream"
    )

# =========================
# HEALTH
# =========================

@app.get("/")
def root():

    return {
        "status": "Fast Coder Running 🚀"
    }

# =========================
# RUN
# =========================

if __name__ == "__main__":

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