import os
import json
from threading import Thread
from fastapi import FastAPI
from fastapi.responses import HTMLResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from typing import List, Optional
from huggingface_hub import hf_hub_download

# ── Model setup ──────────────────────────────────────────────────────────────

MODEL_REPO = "OBLITERATUS/gemma-4-E4B-it-OBLITERATED"
MODEL_FILE = "gemma-4-E4B-it-OBLITERATED-Q4_K_M.gguf"
MAX_CONTEXT_MESSAGES = 4
MAX_NEW_TOKENS = 512

SYSTEM_PROMPT = "You are a helpful AI assistant. Respond to the user's input."

# Download the GGUF model (~4.9 GB, cached after first download)
print(f"Downloading {MODEL_FILE}…")
model_path = hf_hub_download(
    repo_id=MODEL_REPO,
    filename=MODEL_FILE,
    cache_dir=".",
)
print(f"Model downloaded: {model_path}")
print(f"File size: {os.path.getsize(model_path) / (1024*1024):.2f} MB")

# Import llama_cpp after ensuring it's installed
import llama_cpp
from llama_cpp import Llama
print(f"llama-cpp-python version: {llama_cpp.__version__}")

# CPU thread count — match HF free tier (2 vCPUs)
N_THREADS = int(os.environ.get("N_THREADS", "2"))

print("Loading model into memory…")
llm = Llama(
    model_path=model_path,
    n_ctx=2048,           # context window
    n_threads=N_THREADS,
    n_threads_batch=N_THREADS,
    n_gpu_layers=0,       # CPU only on free tier
    verbose=True,         # Enable verbose for better debugging
    use_mmap=False,       # Disable mmap for better compatibility with /tmp or network drives
    use_mlock=False,      # don't lock in RAM
)
print("Model ready ✓")

# ── FastAPI app ──────────────────────────────────────────────────────────────

app = FastAPI()
app.mount("/static", StaticFiles(directory="static"), name="static")


class Message(BaseModel):
    role: str
    content: str

class ChatRequest(BaseModel):
    messages: List[Message]
    memories: Optional[str] = None


@app.get("/", response_class=HTMLResponse)
async def index():
    with open("static/index.html", "r") as f:
        return HTMLResponse(content=f.read())


def build_messages(req_messages, memories=None):
    """Build the message list with system prompt, memories, and truncated history."""
    msgs = []

    # System prompt
    system = SYSTEM_PROMPT
    if memories and memories.strip():
        system += f"\n\nMemories from past conversations:\n{memories[:300]}"
    msgs.append({"role": "system", "content": system})

    # Truncate to last N messages
    history = req_messages[-MAX_CONTEXT_MESSAGES:]
    for m in history:
        content = m["content"][:1500] if len(m["content"]) > 1500 else m["content"]
        msgs.append({"role": m["role"], "content": content})

    return msgs


@app.post("/api/chat")
async def chat(req: ChatRequest):
    messages = [{"role": m.role, "content": m.content} for m in req.messages]
    chat_msgs = build_messages(messages, req.memories)

    def event_stream():
        response = llm.create_chat_completion(
            messages=chat_msgs,
            stream=True,
            max_tokens=MAX_NEW_TOKENS,
            temperature=0.7,
            top_p=0.9,
            top_k=40,
            repeat_penalty=1.1,
        )
        for chunk in response:
            delta = chunk["choices"][0].get("delta", {})
            token = delta.get("content", "")
            if token:
                yield f"data: {json.dumps({'token': token})}\n\n"
        yield "data: [DONE]\n\n"

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


@app.post("/api/summarize")
async def summarize(req: ChatRequest):
    """Summarize a conversation into compact memory bullets."""
    msgs = req.messages[-6:]
    conversation_text = ""
    for m in msgs:
        role = "User" if m.role == "user" else "Assistant"
        conversation_text += f"{role}: {m.content[:300]}\n"

    summary_msgs = [
        {"role": "user", "content": (
            "Summarize this conversation in 3-5 bullet points. "
            "Focus on key facts and user preferences. Be very concise.\n\n"
            f"{conversation_text}\n\nBullets:"
        )}
    ]

    response = llm.create_chat_completion(
        messages=summary_msgs,
        max_tokens=128,
        temperature=0.3,
    )

    result = response["choices"][0]["message"]["content"]
    return {"summary": result.strip()}