import os
import json
from fastapi import FastAPI
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import BaseModel
from typing import List, Optional
from huggingface_hub import hf_hub_download
from llama_cpp import Llama

REPO_ID = "Qwen/Qwen2.5-Coder-3B-Instruct-GGUF"
FILENAME = "qwen2.5-coder-3b-instruct-q4_k_m.gguf"
MODEL_NAME = "Qwen2.5-Coder-3B-Instruct-GGUF"

print("Downloading GGUF model...")
model_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
print("Model downloaded!")

print("Loading model into Llama.cpp...")
llm = Llama(
    model_path=model_path,
    n_ctx=8192,
    n_batch=512,  # Chunked prefill সচল রাখা হলো
    n_threads=2,
    verbose=False
)
print("Llama.cpp ready!")

app = FastAPI()

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

class ChatCompletionRequest(BaseModel):
    model: Optional[str] = MODEL_NAME
    messages: List[ChatMessage]
    max_tokens: Optional[int] = 1024
    temperature: Optional[float] = 0.7
    stream: Optional[bool] = False

@app.get("/")
def root():
    return {"status": "online", "model": MODEL_NAME}

@app.post("/v1/chat/completions")
async def chat_completions(request: ChatCompletionRequest):
    try:
        raw_messages = [{"role": m.role, "content": m.content} for m in request.messages]

        # 🧠 স্মার্ট ট্রিমিং লজিক: কনটেক্সট ১৭কে থেকে কমিয়ে সেফ জোনে আনা
        system_message = None
        if raw_messages and raw_messages[0]["role"] == "system":
            system_message = raw_messages[0]
            chat_messages = raw_messages[1:]
        else:
            chat_messages = raw_messages

        # ৪০০০ টোকেন ≈ ১৬,০০০ ক্যারেক্টার (সুরক্ষার জন্য লিমিট করা হলো)
        max_chars = 15000 
        trimmed_chat = []
        current_chars = len(system_message["content"]) if system_message else 0

        # শেষ দিক থেকে (সবচেয়ে নতুন মেসেজগুলো) নেওয়া শুরু করবে
        for msg in reversed(chat_messages):
            msg_len = len(msg["content"])
            if current_chars + msg_len < max_chars:
                trimmed_chat.insert(0, msg)
                current_chars += msg_len
            else:
                break

        # ফাইনাল মেসেজ লিস্ট তৈরি
        final_messages = [system_message] + trimmed_chat if system_message else trimmed_chat

        max_tokens = min(request.max_tokens or 1024, 2048)

        # ১. স্ট্রিমিং মোড
        if request.stream:
            def stream_generator():
                chunks = llm.create_chat_completion(
                    messages=final_messages,
                    max_tokens=max_tokens,
                    temperature=request.temperature or 0.7,
                    stream=True
                )
                for chunk in chunks:
                    if "model" in chunk:
                        chunk["model"] = MODEL_NAME
                    yield f"data: {json.dumps(chunk)}\n\n"
                yield "data: [DONE]\n\n"
            
            return StreamingResponse(stream_generator(), media_type="text/event-stream")
        
        # ২. নন-স্ট্রিমিং মোড
        else:
            response = llm.create_chat_completion(
                messages=final_messages,
                max_tokens=max_tokens,
                temperature=request.temperature or 0.7,
                stream=False
            )
            response["model"] = MODEL_NAME
            return response

    except Exception as e:
        return JSONResponse(
            status_code=500,
            content={"error": str(e)}
        )