import os
import uvicorn
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List, Dict
import gradio as gr
from huggingface_hub import hf_hub_download
from llama_cpp import Llama

# ==========================================
# 1. Model Setup & Quantization Trick (4-bit)
# ==========================================
# Using the highly reliable bartowski repository for the GGUF file
repo_id = "bartowski/Qwen2.5-7B-Instruct-GGUF"
filename = "Qwen2.5-7B-Instruct-Q4_K_M.gguf"

print("Downloading/Verifying 4-bit GGUF model...")
model_path = hf_hub_download(repo_id=repo_id, filename=filename)

# Load Llama engine (fits in ~4GB RAM)
llm = Llama(
    model_path=model_path,
    n_ctx=2048,       # Context window (memory limit)
    n_threads=2,      # HF Free Spaces have 2 vCPUs
    verbose=False
)

# ==========================================
# 2. FastAPI Setup (API Access)
# ==========================================
app = FastAPI(title="Qwen 7B API & UI Space")

class ChatRequest(BaseModel):
    messages: List[Dict[str, str]] # e.g., [{"role": "user", "content": "hi"}]
    max_tokens: int = 512
    temperature: float = 0.7

@app.post("/api/chat")
async def api_chat(request: ChatRequest):
    """
    API Endpoint for external access.
    Expects ChatML format memory/history.
    """
    try:
        # Build prompt from memory/messages
        prompt = "<|im_start|>system\nYou are a helpful AI assistant.<|im_end|>\n"
        for msg in request.messages:
            role = msg.get("role", "user")
            content = msg.get("content", "")
            prompt += f"<|im_start|>{role}\n{content}<|im_end|>\n"
        
        prompt += "<|im_start|>assistant\n"

        # Generate response
        output = llm(
            prompt,
            max_tokens=request.max_tokens,
            stop=["<|im_end|>", "<|im_start|>"],
            temperature=request.temperature,
            top_p=0.9
        )
        
        return {
            "reply": output['choices'][0]['text'].strip(),
            "usage": output.get('usage', {})
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# ==========================================
# 3. Gradio UI Setup (With Session Memory)
# ==========================================
def chat_stream(message, history):
    """
    Handles the Gradio UI chat. 'history' acts as session memory.
    """
    prompt = "<|im_start|>system\nYou are a helpful AI assistant.<|im_end|>\n"
    
    # Inject session memory
    for user_msg, bot_msg in history:
        prompt += f"<|im_start|>user\n{user_msg}<|im_end|>\n"
        prompt += f"<|im_start|>assistant\n{bot_msg}<|im_end|>\n"
        
    # Current message
    prompt += f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
    
    # Stream output to UI
    stream = llm(
        prompt,
        max_tokens=512,
        stop=["<|im_end|>", "<|im_start|>"],
        stream=True,
        temperature=0.7
    )
    
    response = ""
    for chunk in stream:
        token = chunk['choices'][0]['text']
        response += token
        yield response

# Create the Gradio interface (Removed incompatible theme/layout arguments)
ui = gr.ChatInterface(
    fn=chat_stream,
    title="Qwen 2.5 7B (4-bit GGUF)",
    description="UI features per-user session memory. API is available at `/api/chat`."
)

# ==========================================
# 4. Mount UI onto API & Run
# ==========================================
# This serves the Gradio UI at the root URL "/"
app = gr.mount_gradio_app(app, ui, path="/")

if __name__ == "__main__":
    # Run the Uvicorn server on port 7860 (required by Hugging Face Spaces)
    uvicorn.run(app, host="0.0.0.0", port=7860)