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

# Download model from Unsloth repository
model_path = hf_hub_download(
    repo_id="unsloth/DeepSeek-R1-Distill-Qwen-7B-GGUF",
    filename="DeepSeek-R1-Distill-Qwen-7B-Q4_K_M.gguf"
)

# Instantiate Llama model
llm = Llama(
    model_path=model_path,
    n_ctx=4096,
    n_threads=2
)

app = FastAPI(title="Lumora-7b DeepSeek-R1 Streaming API")

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

class ChatRequest(BaseModel):
    messages: List[Message]
    temperature: Optional[float] = 0.6
    max_tokens: Optional[int] = 1024

async def response_generator(prompt: str, temperature: float, max_tokens: int):
    # Call llama-cpp-python generator loop directly with stream=True
    response_stream = llm(
        prompt,
        max_tokens=max_tokens,
        temperature=temperature,
        stream=True,
        stop=["<|im_end|>", "<|im_start|>"]
    )
    
    in_think_block = False
    
    for chunk in response_stream:
        token = chunk["choices"][0]["text"]
        if not token:
            continue
            
        delta = {}
        
        # Track active block transitions to accurately isolate the reasoning layers
        if "<think>" in token:
            in_think_block = True
            # Strip the structural tag out of the reasoning output payload
            token = token.replace("<think>", "")
            if token:
                delta["reasoning_content"] = token
        elif "</think>" in token:
            in_think_block = False
            token = token.replace("</think>", "")
            if token:
                delta["content"] = token
        else:
            if in_think_block:
                delta["reasoning_content"] = token
            else:
                delta["content"] = token

        # Wrap structural data matching openAI delta layout formatting specs
        if delta:
            json_data = {
                "choices": [
                    {
                        "delta": delta,
                        "finish_reason": None
                    }
                ]
            }
            # SSE streams mandate prefixing with 'data: ' and ending with two newlines
            yield f"data: {json.dumps(json_data)}\n\n"
            
    # Yield structural finalization token string to close pipeline cleanly
    final_json = {"choices": [{"delta": {}, "finish_reason": "stop"}]}
    yield f"data: {json.dumps(final_json)}\n\n"
    yield "data: [DONE]\n\n"

@app.post("/v1/chat/completions")
async def chat_completions(request: ChatRequest):
    try:
        # Loop over past history elements to build out system context bounds
        prompt = ""
        for msg in request.messages:
            prompt += f"<|im_start|>{msg.role}\n{msg.content}<|im_end|>\n"
        prompt += "<|im_start|>assistant\n"

        return StreamingResponse(
            response_generator(prompt, request.temperature, request.max_tokens),
            media_type="text/event-stream",
            headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)