from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List, Optional
from huggingface_hub import hf_hub_download
from llama_cpp import Llama

app = FastAPI(
    title="Gemma 2 API",
    description="Optimoitu Gemma API 2 vCPU / 16GB RAM ympäristölle"
)

# Ladataan malli Hugging Facesta (Gemma 2 2B Instruct - 4-bit quant)
REPO_ID = "bartowski/gemma-2-2b-it-GGUF"
FILENAME = "gemma-2-2b-it-Q4_K_M.gguf"

print("Ladataan mallitiedostoa...")
model_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
print(f"Malli ladattu osoitteeseen: {model_path}")

# Alustetaan llama-cpp hyödyntämään molempia vCPU-ytimiä
llm = Llama(
    model_path=model_path,
    n_ctx=4096,      # Konteksti-ikkuna
    n_threads=2,     # 2 vCPU
    n_batch=512,
    verbose=False
)

# Pyyntömallit
class Message(BaseModel):
    role: str
    content: str

class ChatRequest(BaseModel):
    messages: List[Message]
    max_tokens: Optional[int] = 512
    temperature: Optional[float] = 0.7
    top_p: Optional[float] = 0.9

class PromptRequest(BaseModel):
    prompt: str
    max_tokens: Optional[int] = 512
    temperature: Optional[float] = 0.7

@app.get("/")
def root():
    return {
        "status": "online",
        "model": "Gemma-2-2B-IT-Q4_K_M",
        "endpoints": ["/v1/chat/completions", "/generate", "/docs"]
    }

# 1. Yksinkertainen Prompt API
@app.post("/generate")
def generate(req: PromptRequest):
    try:
        output = llm(
            req.prompt,
            max_tokens=req.max_tokens,
            temperature=req.temperature,
            stop=["<end_of_turn>", "<eos>"]
        )
        return {"response": output["choices"][0]["text"]}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# 2. OpenAI-yhteensopiva Chat Completions API
@app.post("/v1/chat/completions")
def chat_completions(req: ChatRequest):
    try:
        # Muodostetaan Gemma 2 -spesifinen prompt-formaatti
        formatted_prompt = ""
        for msg in req.messages:
            formatted_prompt += f"<start_of_turn>{msg.role}\n{msg.content}<end_of_turn>\n"
        formatted_prompt += "<start_of_turn>model\n"

        output = llm(
            formatted_prompt,
            max_tokens=req.max_tokens,
            temperature=req.temperature,
            top_p=req.top_p,
            stop=["<end_of_turn>", "<eos>", "<start_of_turn>"]
        )

        response_text = output["choices"][0]["text"].strip()

        return {
            "id": output.get("id", "chatcmpl-gemma"),
            "object": "chat.completion",
            "choices": [
                {
                    "index": 0,
                    "message": {
                        "role": "assistant",
                        "content": response_text
                    },
                    "finish_reason": output["choices"][0].get("finish_reason", "stop")
                }
            ],
            "usage": output.get("usage", {})
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))