"""llamafile provisioning — single-binary local LLM.

Downloads and manages a llamafile binary that provides an OpenAI-compatible
API on localhost without Docker, K3D, or any container runtime.

Mozilla project, Apache 2.0 license.
"""

from __future__ import annotations

import os
import socket
import stat
import subprocess
from pathlib import Path

# Bonsai 1.7B GGUF — small enough for CPU, good enough for basic Q&A
LLAMAFILE_URL = "https://huggingface.co/Mozilla/llamafile/resolve/main/llamafile-0.8.17"
MODEL_URL = "https://huggingface.co/prism-ml/Bonsai-1.7B-gguf/resolve/main/Bonsai-1.7B.gguf"
MODEL_NAME = "Bonsai-1.7B.gguf"
DEFAULT_PORT = 8080


def get_llamafile_dir() -> Path:
    """Return ~/.axi/llamafile/, creating it if needed."""
    d = Path.home() / ".axi" / "llamafile"
    d.mkdir(parents=True, exist_ok=True)
    return d


def is_llamafile_installed() -> bool:
    """Check if llamafile binary and model are downloaded."""
    d = get_llamafile_dir()
    return (d / "llamafile").exists() and (d / MODEL_NAME).exists()


def is_llamafile_running(port: int = DEFAULT_PORT) -> bool:
    """Check if llamafile server is responding on *port*."""
    try:
        with socket.create_connection(("localhost", port), timeout=1):
            return True
    except (ConnectionRefusedError, OSError):
        return False


def download_llamafile(progress_callback=None) -> Path:
    """Download llamafile binary if not present."""
    d = get_llamafile_dir()
    binary = d / "llamafile"

    if binary.exists():
        return binary

    import requests

    if progress_callback:
        progress_callback("Downloading llamafile runtime...")

    resp = requests.get(LLAMAFILE_URL, stream=True, timeout=30)
    resp.raise_for_status()

    with open(binary, "wb") as f:
        for chunk in resp.iter_content(chunk_size=8192):
            f.write(chunk)

    # Make executable
    binary.chmod(binary.stat().st_mode | stat.S_IEXEC | stat.S_IXGRP | stat.S_IXOTH)

    return binary


def download_model(progress_callback=None) -> Path:
    """Download Bonsai model if not present."""
    d = get_llamafile_dir()
    model = d / MODEL_NAME

    if model.exists():
        return model

    import requests

    if progress_callback:
        progress_callback(f"Downloading {MODEL_NAME} (1.7GB)...")

    resp = requests.get(MODEL_URL, stream=True, timeout=30)
    resp.raise_for_status()

    total = int(resp.headers.get("content-length", 0))
    downloaded = 0

    with open(model, "wb") as f:
        for chunk in resp.iter_content(chunk_size=65536):
            f.write(chunk)
            downloaded += len(chunk)
            if progress_callback and total:
                pct = int(downloaded / total * 100)
                progress_callback(f"Downloading {MODEL_NAME}... {pct}%")

    return model


def start_llamafile(port: int = DEFAULT_PORT, background: bool = True) -> bool:
    """Start llamafile server.

    Returns True if started successfully.
    """
    if is_llamafile_running(port):
        return True

    d = get_llamafile_dir()
    binary = d / "llamafile"
    model = d / MODEL_NAME

    if not binary.exists() or not model.exists():
        return False

    cmd = [
        str(binary),
        "-m",
        str(model),
        "--port",
        str(port),
        "--host",
        "127.0.0.1",
        "-c",
        "4096",
        "-t",
        "2",
        "--log-disable",
    ]

    if background:
        # Start as background process
        pid_file = d / "llamafile.pid"
        proc = subprocess.Popen(
            cmd,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
            start_new_session=True,
        )
        pid_file.write_text(str(proc.pid))

        # Wait for server to be ready
        import time

        for _ in range(30):  # 30 second timeout
            time.sleep(1)
            if is_llamafile_running(port):
                return True
        return False
    else:
        subprocess.run(cmd, check=True)
        return True


def stop_llamafile() -> bool:
    """Stop the llamafile server."""
    d = get_llamafile_dir()
    pid_file = d / "llamafile.pid"

    if not pid_file.exists():
        return False

    try:
        pid = int(pid_file.read_text().strip())
        os.kill(pid, 15)  # SIGTERM
        pid_file.unlink()
        return True
    except (ProcessLookupError, ValueError):
        pid_file.unlink(missing_ok=True)
        return False


def provision(progress_callback=None) -> dict:
    """Full provisioning: download + start.

    Returns status dict.
    """
    result = {"binary": False, "model": False, "running": False, "port": DEFAULT_PORT}

    try:
        download_llamafile(progress_callback)
        result["binary"] = True
    except Exception as e:
        result["error"] = f"Failed to download llamafile: {e}"
        return result

    try:
        download_model(progress_callback)
        result["model"] = True
    except Exception as e:
        result["error"] = f"Failed to download model: {e}"
        return result

    try:
        started = start_llamafile()
        result["running"] = started
        if not started:
            result["error"] = "llamafile started but server not responding"
    except Exception as e:
        result["error"] = f"Failed to start llamafile: {e}"

    return result


def get_status() -> dict:
    """Get current llamafile status."""
    return {
        "installed": is_llamafile_installed(),
        "running": is_llamafile_running(),
        "port": DEFAULT_PORT,
        "model": MODEL_NAME,
        "path": str(get_llamafile_dir()),
    }
