# ruff: noqa: T201
"""Configuration import handler for Llama-cpp."""

import json
import os
import struct
from pathlib import Path
from typing import Literal

import gguf
from numpy import memmap
from platformdirs import user_data_dir

from yera.config2.dataclasses import (
    CredentialsMap,
    LLMConfig,
    ModelRegistry,
)
from yera.config2.keyring import DevKeyring
from yera.config2.setup_handlers.base import BaseProviderSetup


def _decode_value(value: memmap | bytes) -> str | int:
    if hasattr(value, "tobytes"):
        value = value.tobytes()

    if isinstance(value, bytes):
        # Try to decode as UTF-8 string first
        try:
            decoded = value.decode("utf-8")
            # If it's printable and non-empty, prefer string interpretation
            if decoded.isprintable() and decoded.strip():
                return decoded
        except UnicodeDecodeError:
            pass

        # For 4-byte values that aren't valid strings, try integer unpacking
        if len(value) == 4:
            try:
                return struct.unpack("<I", value)[0]
            except struct.error:
                pass

    return "[could not decode]"


def _get_context_length(meta: dict) -> int | str:
    keys = [k for k in meta if k.endswith(".context_length")]
    if len(keys) == 1:
        return meta[keys[0]]
    return "Unknown"


def _find_project_root(
    srcs: tuple[str, ...] | None = None,
    markers: tuple[str, ...] = ("pyproject.toml", ".git"),
) -> Path:
    if not srcs:
        srcs = (str(Path.cwd()),)

    path = Path(srcs[0]).resolve()

    for parent in [path, *path.parents]:
        if any((parent / marker).exists() for marker in markers):
            return parent

    return path


def _get_models_dir(
    location: Literal["local", "global"],
) -> Path:
    if location == "local":
        return _find_project_root() / "models"
    if location == "global":
        return Path(user_data_dir("yera")) / "models"
    raise ValueError(f"Unknown location: {location}")


class LlamaCppSetup(BaseProviderSetup):
    """Setup handler for llama-cpp.

    Manages finding the model directory, finding ggufs, extracting their
    metadata and creating a catalogue json that maps IDs to local files.
    """

    def __init__(self):
        """Initialise the LlamaCpp setup handler."""
        super().__init__("llama-cpp")

    def detect_creds(self) -> dict[str, str] | None:
        """Detect LlamaCpp models directory.

        The priority is
            1. YERA_MODELS_DIR env variable
            2. ./models
            3. ~/.local/yera/models

        Returns:
            Dictionary containing the model directory path, otherwise
            None if no model directory is found,

        """
        # 1. Check environment variable first
        if env_dir := os.getenv("YERA_MODEL_DIR"):
            return {"models_dir": env_dir}

        # 2. Try local project directory
        local_dir = _get_models_dir("local")
        if local_dir.exists():
            return {"models_dir": str(local_dir)}

        # 3. Fall back to global user directory
        global_dir = _get_models_dir("global")
        return {"models_dir": str(global_dir)}

    def ask_for_creds(self) -> dict[str, str] | None:
        """Prompt the user to enter their models directory manually.

        Returns:
            Dictionary containing the model directory path, otherwise
            None if user enters an invalid path.

        """
        print("Please enter the path to your models directory.")
        models_dir = Path(input(" > "))
        return {"models_dir": str(models_dir)}

    def validate_creds(self, creds: dict[str, str]) -> None:
        """Validate the path to the models directory.

        Args:
            creds: Dictionary containing the model directory path.

        Raises:
            FileNotFoundError: Raised if the path does not exist.

        """
        models_dir = Path(creds["models_dir"])
        if not models_dir.exists():
            raise FileNotFoundError(models_dir)

    def _make_model_cfg(self, model: dict[str, str]) -> LLMConfig:
        provider = model.get("organization", "unknown")
        gguf_path = Path(model["gguf_path"])

        model_id = gguf_path.stem.replace(".", "p")
        return LLMConfig(
            id=f"llamacpp.{provider}.{model_id}".lower(),
            display_name=model["name"],
            credentials=self.provider_name,
            provider=provider,
            interface="llama-cpp",
            model_id=model_id,
        )

    def fetch_models(self, creds_map: CredentialsMap) -> ModelRegistry:
        """Fetch information from ggufs in models directory and build registry.

        Args:
            creds_map: Dictionary containing the model directory path.

        Returns:
            ModelRegistry containing configurations for all the models in the
            models directory.

        """
        creds = {
            k: DevKeyring.get(f"providers.{v}")
            for k, v in creds_map.providers[self.provider_name].items()
        }

        models_dir = Path(creds["models_dir"])

        catalogue = {}

        model_configs = []
        for gguf_path in list(models_dir.rglob("*.gguf")):
            reader = gguf.GGUFReader(gguf_path)

            data = {"gguf_path": str(gguf_path)}
            other = {}

            for k, v in reader.fields.items():
                val = v.parts[v.data[0]]
                if k.startswith("general."):
                    data[v.name[8:]] = _decode_value(val)
                else:
                    other[v.name] = _decode_value(val)

            data["quant"] = gguf_path.stem.split("-")[-1]
            data["context_length"] = _get_context_length(other)

            model_cfg = self._make_model_cfg(data)
            catalogue[model_cfg.id] = str(gguf_path)
            model_configs.append(model_cfg)

        (models_dir / "catalogue.json").write_text(json.dumps(catalogue, indent=2))

        names = "  \n   * ".join(m.display_name for m in model_configs)
        print(f"  Adding {len(model_configs)} LLMs from {models_dir}:\n   * {names}")

        return ModelRegistry(
            llm={m.id: m for m in model_configs},
        )
