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

from fastapi import FastAPI, Request, HTTPException, Depends
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, StreamingResponse
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials

from huggingface_hub import hf_hub_download
from llama_cpp import Llama


# ============================================================
# CONFIG
# ============================================================

APP_NAME = "Varsh LLM"

MODEL_ID = "qwen2.5-coder-7b-instruct-q4_k_m"

MODEL_REPO = "Qwen/Qwen2.5-Coder-7B-Instruct-GGUF"
MODEL_FILE = "qwen2.5-coder-7b-instruct-q4_k_m.gguf"

HOST = "0.0.0.0"
PORT = int(os.getenv("PORT", "7860"))

CUSTOM_API_KEY = os.getenv("CUSTOM_API_KEY")

if not CUSTOM_API_KEY:
    raise RuntimeError(
        "CUSTOM_API_KEY is not configured. "
        "Add CUSTOM_API_KEY in Hugging Face Space Settings > Secrets."
    )


# ============================================================
# CONTEXT SETTINGS
# ============================================================

# 16 GB RAM + 7B Q4 CPU deployment.
# 8192 is a safer target than 16384.
CONTEXT_SIZE = 8192

# Reserve room for the model's answer.
RESERVED_OUTPUT_TOKENS = 1024

# Maximum amount of input context we allow.
MAX_INPUT_TOKENS = (
    CONTEXT_SIZE - RESERVED_OUTPUT_TOKENS
)


# ============================================================
# FASTAPI
# ============================================================

app = FastAPI(
    title=APP_NAME,
    version="4.0.0",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

security = HTTPBearer(auto_error=False)


# ============================================================
# AUTH
# ============================================================

def verify_api_key(
    credentials: HTTPAuthorizationCredentials =
    Depends(security),
):
    if credentials is None:
        raise HTTPException(
            status_code=401,
            detail="Missing Authorization header.",
        )

    if credentials.scheme.lower() != "bearer":
        raise HTTPException(
            status_code=401,
            detail="Authorization must use Bearer authentication.",
        )

    if credentials.credentials != CUSTOM_API_KEY:
        raise HTTPException(
            status_code=401,
            detail="Invalid API key.",
        )

    return True


# ============================================================
# MODEL
# ============================================================

print("=" * 70)
print("Starting Varsh LLM")
print("=" * 70)

print(f"Model repository: {MODEL_REPO}")
print(f"Model file:       {MODEL_FILE}")
print(f"Context size:     {CONTEXT_SIZE}")
print(f"Max input:        {MAX_INPUT_TOKENS}")
print(f"Reserved output:  {RESERVED_OUTPUT_TOKENS}")

print("Checking/downloading model...")

model_path = hf_hub_download(
    repo_id=MODEL_REPO,
    filename=MODEL_FILE,
)

print(f"Model path: {model_path}")
print("Loading Qwen 2.5 Coder 7B...")

llm = Llama(
    model_path=model_path,

    # CPU deployment
    n_ctx=CONTEXT_SIZE,

    n_threads=2,
    n_threads_batch=2,

    # Moderate batch for 2 vCPU.
    n_batch=256,

    use_mmap=True,
    use_mlock=False,

    verbose=False,
)

print("Model loaded successfully.")

print("=" * 70)
print("VARSH LLM API READY")
print("=" * 70)
print(f"Listening on:       {HOST}:{PORT}")
print("Root endpoint:      /")
print("Models endpoint:    /v1/models")
print("Chat endpoint:      /v1/chat/completions")
print("Authentication:     Bearer API key")
print("Streaming:          enabled")
print(f"Context:            {CONTEXT_SIZE}")
print("=" * 70)


# llama.cpp generation is serialized because this is
# a 2-vCPU CPU deployment.
generation_lock = asyncio.Lock()


# ============================================================
# CONTENT NORMALIZATION
# ============================================================

def content_to_text(content: Any) -> str:

    if content is None:
        return ""

    if isinstance(content, str):
        return content

    if isinstance(content, list):

        parts = []

        for item in content:

            if isinstance(item, str):
                parts.append(item)
                continue

            if not isinstance(item, dict):
                continue

            item_type = item.get("type")

            if item_type == "text":

                text = item.get("text", "")

                if isinstance(text, str):
                    parts.append(text)

                continue

            # Some clients use:
            # {"text": "..."}
            if isinstance(
                item.get("text"),
                str,
            ):
                parts.append(
                    item["text"]
                )

        return "\n".join(parts)

    if isinstance(content, dict):

        if isinstance(
            content.get("text"),
            str,
        ):
            return content["text"]

        if isinstance(
            content.get("content"),
            str,
        ):
            return content["content"]

        return json.dumps(
            content,
            ensure_ascii=False,
        )

    return str(content)


# ============================================================
# MESSAGE NORMALIZATION
# ============================================================

def normalize_messages(messages: Any) -> list:

    if not isinstance(messages, list):
        raise HTTPException(
            status_code=400,
            detail="'messages' must be an array.",
        )

    result = []

    for message in messages:

        if not isinstance(message, dict):
            continue

        role = message.get(
            "role",
            "user",
        )

        if not isinstance(role, str):
            role = "user"

        content = content_to_text(
            message.get(
                "content",
                "",
            )
        )

        normalized = {
            "role": role,
            "content": content,
        }

        if isinstance(
            message.get("name"),
            str,
        ):
            normalized["name"] = (
                message["name"]
            )

        if isinstance(
            message.get("tool_call_id"),
            str,
        ):
            normalized["tool_call_id"] = (
                message["tool_call_id"]
            )

        result.append(normalized)

    if not result:
        raise HTTPException(
            status_code=400,
            detail="'messages' cannot be empty.",
        )

    return result


# ============================================================
# TOKEN COUNTING
# ============================================================

def count_message_tokens(messages: list) -> int:
    """
    Use llama.cpp's actual tokenizer instead of estimating
    tokens using character length.
    """

    try:

        # llama.cpp's chat template will ultimately create
        # a prompt. We use a conservative approximation here.
        total = 0

        for message in messages:

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

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

            text = (
                f"<|im_start|>{role}\n"
                f"{content}"
                f"<|im_end|>\n"
            )

            tokens = llm.tokenize(
                text.encode(
                    "utf-8"
                ),
                add_bos=False,
            )

            total += len(tokens)

        # Assistant generation marker.
        assistant_tokens = llm.tokenize(
            b"<|im_start|>assistant\n",
            add_bos=False,
        )

        total += len(
            assistant_tokens
        )

        return total

    except Exception:

        # Safe fallback.
        characters = sum(
            len(
                str(
                    m.get(
                        "content",
                        "",
                    )
                )
            )
            for m in messages
        )

        return max(
            1,
            characters // 4,
        )


# ============================================================
# CONTEXT MANAGEMENT
# ============================================================

def fit_messages_to_context(
    messages: list,
    max_tokens: int,
) -> tuple[list, int]:

    if not messages:
        return messages, 0

    current_tokens = count_message_tokens(
        messages
    )

    if current_tokens <= max_tokens:
        return messages, current_tokens

    print(
        f"Context too large: "
        f"{current_tokens} tokens > "
        f"{max_tokens}. "
        f"Trimming old messages."
    )

    # --------------------------------------------------------
    # Preserve system messages.
    # --------------------------------------------------------

    system_messages = [
        m
        for m in messages
        if m.get("role") == "system"
    ]

    non_system_messages = [
        m
        for m in messages
        if m.get("role") != "system"
    ]

    # --------------------------------------------------------
    # Keep newest messages first.
    # --------------------------------------------------------

    kept = list(system_messages)

    # Add newest conversation messages backwards.
    selected = []

    for message in reversed(
        non_system_messages
    ):

        candidate = (
            system_messages
            + list(reversed(
                selected + [message]
            ))
        )

        tokens = count_message_tokens(
            candidate
        )

        if tokens <= max_tokens:
            selected.append(message)
        else:
            break

    selected.reverse()

    final_messages = (
        system_messages
        + selected
    )

    final_tokens = count_message_tokens(
        final_messages
    )

    # --------------------------------------------------------
    # If system prompt itself is too large,
    # truncate the system content.
    # --------------------------------------------------------

    if final_tokens > max_tokens:

        print(
            "System context is too large. "
            "Truncating system message."
        )

        new_messages = []

        remaining = max_tokens

        for message in final_messages:

            if (
                message.get("role")
                != "system"
            ):
                continue

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

            # Tokenize system content.
            tokens = llm.tokenize(
                content.encode(
                    "utf-8"
                ),
                add_bos=False,
            )

            if len(tokens) <= remaining:

                new_messages.append(
                    message
                )

                remaining -= len(tokens)

            else:

                # Approximate character cutoff.
                ratio = (
                    remaining
                    / max(
                        len(tokens),
                        1,
                    )
                )

                cutoff = max(
                    1,
                    int(
                        len(content)
                        * ratio
                    ),
                )

                truncated = content[
                    :cutoff
                ]

                new_messages.append(
                    {
                        "role": "system",
                        "content": truncated,
                    }
                )

                remaining = 0
                break

        # Add as many newest messages as fit.
        for message in selected:

            candidate = (
                new_messages
                + [message]
            )

            if (
                count_message_tokens(
                    candidate
                )
                <= max_tokens
            ):
                new_messages.append(
                    message
                )

        final_messages = new_messages

        final_tokens = count_message_tokens(
            final_messages
        )

    print(
        f"Context after trimming: "
        f"{final_tokens} tokens"
    )

    return (
        final_messages,
        final_tokens,
    )


# ============================================================
# HELPERS
# ============================================================

def make_id():
    return f"chatcmpl-{uuid.uuid4().hex}"


def safe_int(
    value,
    default,
):
    try:
        return int(value)
    except Exception:
        return default


def safe_float(
    value,
    default,
):
    try:
        return float(value)
    except Exception:
        return default


# ============================================================
# ROOT
# ============================================================

@app.get("/")
async def root():

    return {
        "status": "healthy",
        "service": APP_NAME,
        "security": "enabled",
        "model": MODEL_ID,
        "endpoint": "/v1/chat/completions",
        "models_endpoint": "/v1/models",
        "streaming": True,
        "context_size": CONTEXT_SIZE,
    }


# ============================================================
# HEALTH
# ============================================================

@app.get("/health")
async def health():

    return {
        "status": "ok",
        "model": MODEL_ID,
        "context_size": CONTEXT_SIZE,
    }


# ============================================================
# MODELS
# ============================================================

@app.get(
    "/v1/models",
    dependencies=[Depends(verify_api_key)],
)
async def models():

    return {
        "object": "list",
        "data": [
            {
                "id": MODEL_ID,
                "object": "model",
                "created": int(time.time()),
                "owned_by": "varsshancoder",
            }
        ],
    }


# ============================================================
# CHAT
# ============================================================

@app.post(
    "/v1/chat/completions",
    dependencies=[Depends(verify_api_key)],
)
async def chat_completions(
    request: Request,
):

    # --------------------------------------------------------
    # JSON
    # --------------------------------------------------------

    try:
        body = await request.json()

    except Exception:
        raise HTTPException(
            status_code=400,
            detail="Invalid JSON request body.",
        )

    if not isinstance(body, dict):
        raise HTTPException(
            status_code=400,
            detail="Request body must be a JSON object.",
        )

    # --------------------------------------------------------
    # Messages
    # --------------------------------------------------------

    messages = normalize_messages(
        body.get(
            "messages",
            [],
        )
    )

    # --------------------------------------------------------
    # Model
    # --------------------------------------------------------

    requested_model = body.get(
        "model",
        MODEL_ID,
    )

    if requested_model == f"varsh/{MODEL_ID}":
        requested_model = MODEL_ID

    # --------------------------------------------------------
    # Parameters
    # --------------------------------------------------------

    temperature = safe_float(
        body.get(
            "temperature",
            0.2,
        ),
        0.2,
    )

    top_p = safe_float(
        body.get(
            "top_p",
            0.95,
        ),
        0.95,
    )

    max_tokens = safe_int(
        body.get(
            "max_tokens",
            512,
        ),
        512,
    )

    stream = bool(
        body.get(
            "stream",
            False,
        )
    )

    temperature = max(
        0.0,
        min(
            temperature,
            2.0,
        ),
    )

    top_p = max(
        0.01,
        min(
            top_p,
            1.0,
        ),
    )

    # --------------------------------------------------------
    # Context management
    # --------------------------------------------------------

    messages, prompt_tokens = (
        fit_messages_to_context(
            messages,
            MAX_INPUT_TOKENS,
        )
    )

    # Never request more than the remaining context.
    available_output = (
        CONTEXT_SIZE
        - prompt_tokens
        - 16
    )

    if available_output < 1:
        raise HTTPException(
            status_code=400,
            detail=(
                "Prompt is too large even after "
                "context trimming."
            ),
        )

    max_tokens = min(
        max_tokens,
        available_output,
    )

    print(
        f"Request: prompt={prompt_tokens}, "
        f"max_tokens={max_tokens}, "
        f"stream={stream}"
    )

    # ========================================================
    # STREAM
    # ========================================================

    if stream:

        chat_id = make_id()
        created = int(time.time())

        async def stream_generator():

            # ------------------------------------------------
            # Immediate first SSE event.
            # ------------------------------------------------

            first_chunk = {
                "id": chat_id,
                "object": "chat.completion.chunk",
                "created": created,
                "model": requested_model,
                "choices": [
                    {
                        "index": 0,
                        "delta": {
                            "role": "assistant",
                            "content": "",
                        },
                        "finish_reason": None,
                    }
                ],
            }

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

            try:

                async with generation_lock:

                    result = (
                        llm.create_chat_completion(
                            messages=messages,
                            temperature=temperature,
                            top_p=top_p,
                            max_tokens=max_tokens,
                            stream=True,
                        )
                    )

                    for chunk in result:

                        choices = chunk.get(
                            "choices",
                            [],
                        )

                        if not choices:
                            continue

                        choice = choices[0]

                        delta = choice.get(
                            "delta",
                            {},
                        )

                        content = delta.get(
                            "content",
                            "",
                        )

                        if content is None:
                            content = ""

                        output_chunk = {
                            "id": chat_id,
                            "object": "chat.completion.chunk",
                            "created": created,
                            "model": requested_model,
                            "choices": [
                                {
                                    "index": 0,
                                    "delta": {
                                        "content": content,
                                    },
                                    "finish_reason": choice.get(
                                        "finish_reason"
                                    ),
                                }
                            ],
                        }

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

                # ------------------------------------------------
                # Final event.
                # ------------------------------------------------

                final_chunk = {
                    "id": chat_id,
                    "object": "chat.completion.chunk",
                    "created": created,
                    "model": requested_model,
                    "choices": [
                        {
                            "index": 0,
                            "delta": {},
                            "finish_reason": "stop",
                        }
                    ],
                }

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

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

            except Exception as exc:

                print(
                    "STREAM ERROR:",
                    repr(exc),
                )

                error_chunk = {
                    "id": chat_id,
                    "object": "chat.completion.chunk",
                    "created": created,
                    "model": requested_model,
                    "choices": [
                        {
                            "index": 0,
                            "delta": {
                                "content": (
                                    "\n[Generation error: "
                                    + str(exc)
                                    + "]"
                                )
                            },
                            "finish_reason": "stop",
                        }
                    ],
                }

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

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

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

    # ========================================================
    # NON-STREAM
    # ========================================================

    try:

        async with generation_lock:

            result = (
                llm.create_chat_completion(
                    messages=messages,
                    temperature=temperature,
                    top_p=top_p,
                    max_tokens=max_tokens,
                    stream=False,
                )
            )

        choices = result.get(
            "choices",
            [],
        )

        if choices:

            choice = choices[0]

            message = choice.get(
                "message",
                {},
            )

            output = message.get(
                "content",
                "",
            )

            finish_reason = choice.get(
                "finish_reason",
                "stop",
            )

        else:

            output = ""
            finish_reason = "stop"

        usage = result.get(
            "usage",
            {},
        )

        return JSONResponse(
            content={
                "id": make_id(),
                "object": "chat.completion",
                "created": int(time.time()),
                "model": requested_model,
                "choices": [
                    {
                        "index": 0,
                        "message": {
                            "role": "assistant",
                            "content": output,
                        },
                        "finish_reason": finish_reason,
                    }
                ],
                "usage": {
                    "prompt_tokens": usage.get(
                        "prompt_tokens",
                        prompt_tokens,
                    ),
                    "completion_tokens": usage.get(
                        "completion_tokens",
                        0,
                    ),
                    "total_tokens": usage.get(
                        "total_tokens",
                        prompt_tokens,
                    ),
                },
            }
        )

    except Exception as exc:

        print(
            "GENERATION ERROR:",
            repr(exc),
        )

        raise HTTPException(
            status_code=500,
            detail=f"Generation failed: {str(exc)}",
        )


# ============================================================
# MAIN
# ============================================================

if __name__ == "__main__":

    import uvicorn

    uvicorn.run(
        app,
        host=HOST,
        port=PORT,
    )