import os
import json
import time
import uuid
import threading
from typing import Any

from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from huggingface_hub import HfApi, hf_hub_download
from llama_cpp import Llama


MODEL_ID = os.getenv("MODEL_ID", "zelapi/lilo-cyber")
MODEL_NAME = os.getenv("MODEL_NAME", "lilo-cyber")
MODEL_FILE = os.getenv("MODEL_FILE", "").strip()

N_CTX = int(os.getenv("N_CTX", "1024"))
N_THREADS = int(os.getenv("N_THREADS", str(os.cpu_count() or 2)))
MAX_TOKENS = int(os.getenv("MAX_TOKENS", "256"))

SYSTEM_PROMPT = (
    "Kamu adalah Lilo Cyber, asisten AI yang membantu pemrograman, "
    "keamanan siber defensif, analisis log, dan DevOps. "
    "Jawab dengan jelas, akurat, dan gunakan bahasa pengguna."
)

api = FastAPI(
    title="Lilo Cyber API",
    description="OpenAI-compatible API for Lilo Cyber",
    version="1.0.0",
)

generation_lock = threading.RLock()
llm = None


def get_model():
    global llm

    if llm is not None:
        return llm

    with generation_lock:
        if llm is not None:
            return llm

        filename = MODEL_FILE

        if not filename:
            files = HfApi().list_repo_files(MODEL_ID)

            candidates = [
                name for name in files
                if name.lower().endswith(".gguf")
            ]

            if not candidates:
                raise RuntimeError(
                    f"Tidak ditemukan file GGUF di {MODEL_ID}"
                )

            preferred = [
                name for name in candidates
                if "Q4_K_M" in name.upper()
            ]

            filename = preferred[0] if preferred else candidates[0]

        model_path = hf_hub_download(
            repo_id=MODEL_ID,
            filename=filename,
        )

        llm = Llama(
            model_path=model_path,
            n_ctx=N_CTX,
            n_threads=max(1, N_THREADS),
            n_batch=128,
            verbose=False,
        )

        return llm


def prepare_messages(messages):
    if not isinstance(messages, list) or not messages:
        raise HTTPException(
            status_code=400,
            detail="messages must be a non-empty array",
        )

    result = [
        {
            "role": "system",
            "content": SYSTEM_PROMPT,
        }
    ]

    for message in messages:
        if not isinstance(message, dict):
            raise HTTPException(
                status_code=400,
                detail="Each message must be an object",
            )

        role = message.get("role")
        content = message.get("content")

        if role == "system":
            continue

        if role not in ("user", "assistant"):
            raise HTTPException(
                status_code=400,
                detail=f"Unsupported role: {role}",
            )

        if not isinstance(content, str):
            raise HTTPException(
                status_code=400,
                detail="Only string message content is supported",
            )

        result.append({
            "role": role,
            "content": content,
        })

    if len(result) == 1:
        raise HTTPException(
            status_code=400,
            detail="At least one user or assistant message is required",
        )

    return result


def completion_parameters(body: dict[str, Any]):
    model = body.get("model", MODEL_NAME)

    if model != MODEL_NAME:
        raise HTTPException(
            status_code=404,
            detail=f"Model not found: {model}",
        )

    messages = prepare_messages(body.get("messages"))
    max_tokens = body.get("max_tokens", MAX_TOKENS)
    temperature = body.get("temperature", 0.7)
    top_p = body.get("top_p", 0.9)

    if (
        not isinstance(max_tokens, int)
        or isinstance(max_tokens, bool)
    ):
        raise HTTPException(
            status_code=400,
            detail="max_tokens must be an integer",
        )

    if not isinstance(temperature, (int, float)):
        raise HTTPException(
            status_code=400,
            detail="temperature must be a number",
        )

    if not isinstance(top_p, (int, float)):
        raise HTTPException(
            status_code=400,
            detail="top_p must be a number",
        )

    return {
        "messages": messages,
        "max_tokens": max(1, min(max_tokens, 512)),
        "temperature": max(0.0, min(float(temperature), 2.0)),
        "top_p": max(0.0, min(float(top_p), 1.0)),
        "repeat_penalty": 1.1,
    }


@api.get("/")
def root():
    return {
        "name": "Lilo Cyber API",
        "model": MODEL_NAME,
        "endpoints": [
            "/health",
            "/v1/models",
            "/v1/chat/completions",
        ],
    }


@api.get("/health")
def health():
    return {
        "status": "ok",
        "model": MODEL_NAME,
        "runtime": "cpu",
        "loaded": llm is not None,
    }


@api.get("/v1/models")
def list_models():
    return {
        "object": "list",
        "data": [
            {
                "id": MODEL_NAME,
                "object": "model",
                "created": int(time.time()),
                "owned_by": "zelapi",
            }
        ],
    }


@api.post("/v1/chat/completions")
def chat_completions(body: dict[str, Any]):
    params = completion_parameters(body)
    request_id = "chatcmpl-" + uuid.uuid4().hex
    created = int(time.time())

    if body.get("stream", False):
        def generate():
            try:
                with generation_lock:
                    model = get_model()

                    yield "data: " + json.dumps({
                        "id": request_id,
                        "object": "chat.completion.chunk",
                        "created": created,
                        "model": MODEL_NAME,
                        "choices": [{
                            "index": 0,
                            "delta": {"role": "assistant"},
                            "finish_reason": None,
                        }],
                    }) + "\n\n"

                    for chunk in model.create_chat_completion(
                        **params,
                        stream=True,
                    ):
                        choice = chunk["choices"][0]

                        payload = {
                            "id": request_id,
                            "object": "chat.completion.chunk",
                            "created": created,
                            "model": MODEL_NAME,
                            "choices": [{
                                "index": 0,
                                "delta": choice.get("delta", {}),
                                "finish_reason": choice.get("finish_reason"),
                            }],
                        }

                        yield (
                            "data: "
                            + json.dumps(payload, ensure_ascii=False)
                            + "\n\n"
                        )

                yield "data: [DONE]\n\n"

            except GeneratorExit:
                raise

            except Exception:
                error = {
                    "error": {
                        "message": "Model inference failed",
                        "type": "server_error",
                    }
                }

                yield "data: " + json.dumps(error) + "\n\n"

        return StreamingResponse(
            generate(),
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "X-Accel-Buffering": "no",
            },
        )

    with generation_lock:
        result = get_model().create_chat_completion(**params)

    result["id"] = request_id
    result["created"] = created
    result["model"] = MODEL_NAME

    return result