from __future__ import annotations

import os
import sys
from pathlib import Path

from huggingface_hub import hf_hub_download, list_repo_files


MODEL_REPO = os.getenv("MODEL_REPO", "mradermacher/Llama3.3-8B-Instruct-Thinking-Heretic-Uncensored-Claude-4.5-Opus-High-Reasoning-i1-GGUF")
QUANT = os.getenv("QUANT", "Q4_K_M")            # used to auto-pick the model file
MODEL_FILE = os.getenv("MODEL_FILE", "auto")     # or an exact filename
MMPROJ_FILE = os.getenv("MMPROJ_FILE", "none")   # text-only model: no projector needed
MMPROJ_PREFER = os.getenv("MMPROJ_PREFER", "Q8_0,f16,F16")  # auto-pick order
MODEL_DIR = Path(os.getenv("MODEL_DIR", "/data/models/llama33-8b-thinking"))

LLAMA_SERVER_BIN = os.getenv("LLAMA_SERVER_BIN", "/opt/llama.cpp/llama-server")
LLAMA_HOST = os.getenv("LLAMA_HOST", "0.0.0.0")
LLAMA_PORT = os.getenv("LLAMA_PORT", "7860")

THREADS = os.getenv("THREADS", "8")
CTX_SIZE = os.getenv("CTX_SIZE", "8192")
BATCH_SIZE = os.getenv("BATCH_SIZE", "default")
UBATCH_SIZE = os.getenv("UBATCH_SIZE", "default")
GPU_LAYERS = os.getenv("GPU_LAYERS", "0")
FLASH_ATTN = os.getenv("FLASH_ATTN", "default")   # on | off | auto | default
CACHE_TYPE_K = os.getenv("CACHE_TYPE_K", "default")
CACHE_TYPE_V = os.getenv("CACHE_TYPE_V", "default")
IMAGE_MAX_TOKENS = os.getenv("IMAGE_MAX_TOKENS", "default")  # only for dynamic-resolution models
CHAT_TEMPLATE_FILE = os.getenv("CHAT_TEMPLATE_FILE", "default")  # default = use template embedded in GGUF

TEMPERATURE = os.getenv("TEMPERATURE", "0.6")
TOP_P = os.getenv("TOP_P", "0.95")
TOP_K = os.getenv("TOP_K", "40")
REPEAT_PENALTY = os.getenv("REPEAT_PENALTY", "1.05")


def log(message: str) -> None:
    print(f"[startup] {message}", flush=True)


def is_set(value: str) -> bool:
    return value.strip().lower() not in {"", "default", "auto", "none", "off"}


def pick_files() -> tuple[str, str | None]:
    """Resolve model + mmproj filenames from the repo listing (or env overrides)."""
    files = []
    if MODEL_FILE.strip().lower() == "auto" or MMPROJ_FILE.strip().lower() == "auto":
        files = [f for f in list_repo_files(MODEL_REPO) if f.lower().endswith(".gguf")]
        log(f"GGUF files in repo: {files}")

    # main model
    if MODEL_FILE.strip().lower() == "auto":
        cands = [f for f in files if "mmproj" not in f.lower() and QUANT.lower() in f.lower()]
        if not cands:
            raise RuntimeError(f"No model file matching QUANT={QUANT}; set MODEL_FILE explicitly.")
        model_file = sorted(cands)[0]
    else:
        model_file = MODEL_FILE

    # projector
    mmproj_file = None
    flag = MMPROJ_FILE.strip().lower()
    if flag == "auto":
        mm = [f for f in files if "mmproj" in f.lower()]
        for pref in [p.strip() for p in MMPROJ_PREFER.split(",") if p.strip()]:
            hit = [f for f in mm if pref.lower() in f.lower()]
            if hit:
                mmproj_file = sorted(hit)[0]
                break
        if not mmproj_file and mm:
            mmproj_file = sorted(mm)[0]
        if not mmproj_file:
            log("WARNING: no mmproj file found; image input will not work")
    elif flag not in {"", "none", "off"}:
        mmproj_file = MMPROJ_FILE

    return model_file, mmproj_file


def fetch(filename: str) -> str:
    MODEL_DIR.mkdir(parents=True, exist_ok=True)
    local = MODEL_DIR / filename
    if local.exists():
        log(f"Using cached file: {local}")
        return str(local)
    log(f"Downloading {MODEL_REPO}/{filename}")
    path = hf_hub_download(repo_id=MODEL_REPO, filename=filename, local_dir=str(MODEL_DIR))
    log(f"Ready: {path}")
    return path


def locate_server() -> str:
    """Find llama-server even if the release tarball layout differs; dump diagnostics if not."""
    import shutil
    import subprocess

    root = Path(LLAMA_SERVER_BIN).parent
    candidates = [Path(LLAMA_SERVER_BIN)]
    for base in {root, Path(os.getenv("LLAMA_DIR", "/opt/llama.cpp"))}:
        if base.exists():
            candidates += list(base.rglob("llama-server"))
    found = next((c for c in candidates if c.is_file()), None) or (
        Path(shutil.which("llama-server")) if shutil.which("llama-server") else None
    )
    if not found:
        log(f"llama-server not found. Contents of {root}:")
        subprocess.run(["ls", "-laR", str(root)], check=False)
        raise FileNotFoundError("llama-server binary missing; see listing above")

    os.chmod(found, 0o755)
    log(f"Using llama-server at {found}")
    # Show missing shared libraries / bad loader, the other common ENOENT cause
    subprocess.run(["ldd", str(found)], check=False)
    os.environ["LD_LIBRARY_PATH"] = f"{found.parent}:{os.environ.get('LD_LIBRARY_PATH', '')}"
    return str(found)


def build_command(model_path: str, mmproj_path: str | None, server_bin: str) -> list[str]:
    cmd = [
        server_bin,
        "-m", model_path,
        "--host", LLAMA_HOST,
        "--port", LLAMA_PORT,
        "--threads", THREADS,
        "--ctx-size", CTX_SIZE,
        "--n-gpu-layers", GPU_LAYERS,
        "--parallel", "1",
        "--cont-batching",
        "--temp", TEMPERATURE,
        "--top-p", TOP_P,
        "--top-k", TOP_K,
        "--repeat-penalty", REPEAT_PENALTY,
    ]

    def add(flag: str, value: str) -> None:
        if is_set(value):
            cmd.extend([flag, value])

    add("--batch-size", BATCH_SIZE)
    add("--ubatch-size", UBATCH_SIZE)
    add("--cache-type-k", CACHE_TYPE_K)
    add("--cache-type-v", CACHE_TYPE_V)
    add("-fa", FLASH_ATTN)
    add("--image-max-tokens", IMAGE_MAX_TOKENS)
    add("--chat-template-file", CHAT_TEMPLATE_FILE)

    if mmproj_path:
        cmd.extend(["--mmproj", mmproj_path])
        if GPU_LAYERS == "0":
            cmd.append("--no-mmproj-offload")  # keep the image encoder on CPU

    return cmd


def main() -> None:
    binary_dir = str(Path(LLAMA_SERVER_BIN).parent)
    existing = os.environ.get("LD_LIBRARY_PATH")
    os.environ["LD_LIBRARY_PATH"] = binary_dir if not existing else f"{binary_dir}:{existing}"

    os.environ.setdefault("OMP_NUM_THREADS", THREADS)
    os.environ.setdefault("OPENBLAS_NUM_THREADS", THREADS)
    os.environ.setdefault("MKL_NUM_THREADS", THREADS)

    model_file, mmproj_file = pick_files()
    model_path = fetch(model_file)
    mmproj_path = fetch(mmproj_file) if mmproj_file else None

    server_bin = locate_server()
    cmd = build_command(model_path, mmproj_path, server_bin)
    log("Starting native llama.cpp web UI")
    log(" ".join(cmd))
    os.execvpe(cmd[0], cmd, os.environ)


if __name__ == "__main__":
    try:
        main()
    except Exception as exc:
        print(f"[fatal] {exc}", file=sys.stderr, flush=True)
        raise