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

# =========================
# CONFIG
# =========================

SUPABASE_URL = os.getenv("SUPABASE_URL")
SUPABASE_KEY = os.getenv("SUPABASE_KEY")

supabase = create_client(
    SUPABASE_URL,
    SUPABASE_KEY
)

stop_flags = {}

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

app = FastAPI()

# =========================
# MODEL CONFIG
# =========================

MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"

print("🚀 Loading Fast Qwen Chat...")

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

# =========================
# 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):
    user_id: str
    conversation_id: str
    messages: list
    temperature: float = 0.2
    stream: bool = True
    branch: bool = False
    parent_id: str | None = None

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

SYSTEM_PROMPT = """
You are a fast and intelligent AI assistant.

Rules:
- Answer naturally
- Avoid repetition
- Do not generate fake conversations
- Do not continue forever
- Finish responses properly
- Be concise but complete
"""

# =========================
# 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()

# =========================
# DB FUNCTIONS
# =========================

def get_messages(cid):

    res = supabase.table("messages") \
        .select("role,content") \
        .eq("conversation_id", cid) \
        .order("created_at") \
        .execute()

    return res.data or []

def save_message(
    cid,
    role,
    content,
    parent_id=None,
    branch_id=None
):

    supabase.table("messages").insert({
        "conversation_id": cid,
        "role": role,
        "content": content,
        "parent_id": parent_id,
        "branch_id": branch_id
    }).execute()

def get_next_branch(parent_id):

    res = supabase.table("messages") \
        .select("branch_id") \
        .eq("parent_id", parent_id) \
        .execute()

    existing = [
        m["branch_id"]
        for m in res.data
        if m.get("branch_id")
    ]

    return max(existing, default=0) + 1

# =========================
# BUILD CHAT TEMPLATE
# =========================

def build_inputs(message, cid):

    # =========================
    # FETCH HISTORY
    # =========================

    history = get_messages(cid)

    # =========================
    # KEEP LAST 2 ONLY
    # =========================

    history = history[-2:]

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

    # =========================
    # ADD HISTORY
    # =========================

    for msg in history:

        messages.append({
            "role": msg["role"],
            "content": msg["content"]
        })

    # =========================
    # ADD CURRENT USER MESSAGE
    # =========================

    messages.append({
        "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):

    last_message = req.messages[-1]["content"]

    inputs = build_inputs(
        last_message,
        req.conversation_id
    )

    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.15,
            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)

    def save_async():

        if not req.branch:

            for msg in req.messages:

                save_message(
                    req.conversation_id,
                    msg["role"],
                    msg["content"]
                )

            save_message(
                req.conversation_id,
                "assistant",
                result
            )

        else:

            branch_id = get_next_branch(
                req.parent_id
            )

            save_message(
                req.conversation_id,
                "assistant",
                result,
                parent_id=req.parent_id,
                branch_id=branch_id
            )

    threading.Thread(
        target=save_async
    ).start()

    return {
        "choices": [
            {
                "message": {
                    "role": "assistant",
                    "content": result
                }
            }
        ],
        "done": True
    }

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

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

    last_message = req.messages[-1]["content"]

    inputs = build_inputs(
        last_message,
        req.conversation_id
    )

    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.15,
        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

                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"

        def save_async():

            if not full_text:
                return

            if not req.branch:

                for msg in req.messages:

                    save_message(
                        req.conversation_id,
                        msg["role"],
                        msg["content"]
                    )

                save_message(
                    req.conversation_id,
                    "assistant",
                    full_text
                )

            else:

                branch_id = get_next_branch(
                    req.parent_id
                )

                save_message(
                    req.conversation_id,
                    "assistant",
                    full_text,
                    parent_id=req.parent_id,
                    branch_id=branch_id
                )

        threading.Thread(
            target=save_async
        ).start()

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

# =========================
# FEEDBACK
# =========================

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

    try:

        supabase.table("messages").update({
            "feedback": data.get("feedback")
        }).eq(
            "id",
            data.get("message_id")
        ).execute()

        return {
            "status": "saved"
        }

    except Exception as e:

        return {
            "error": str(e)
        }

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

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

    return {
        "status": "Fast Qwen Chat Running 🚀"
    }

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

if __name__ == "__main__":

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