"""
Model runner.

Dispatches to:
  - llama-cpp-python  for GGUF models
  - HuggingFace transformers  for everything else
"""

from __future__ import annotations

import sys
from pathlib import Path

from rich.console import Console

console = Console(stderr=True)


# ── GGUF / llama.cpp ─────────────────────────────────────────────────────────

def load_llama_cpp():
    """Return the Llama class, or exit with an install hint."""
    try:
        # pylint: disable-next=import-outside-toplevel
        from llama_cpp import Llama  # type: ignore
        return Llama
    except ImportError:
        console.print(
            "\n  [bold red]llama-cpp-python[/] is not installed.\n\n"
            "  Install it with:\n\n"
            "    [bold cyan]pip install llama-cpp-python[/]\n\n"
            "  For GPU acceleration see:\n"
            "    https://llama-cpp-python.readthedocs.io\n"
        )
        sys.exit(1)


def run_gguf_interactive(  # pylint: disable=too-many-locals
    model_path: Path,
    *,
    n_ctx: int = 128000,
    n_gpu_layers: int = -1,
    max_tokens: int = 128000,
    verbose: bool = False,
    system_prompt: str | None = None,
) -> None:
    """
    Start an interactive chat session with a GGUF model via llama-cpp-python.
    """
    LlamaClass = load_llama_cpp()  # pylint: disable=invalid-name

    from rich.panel import Panel  # pylint: disable=import-outside-toplevel

    console.print("\n  [dim]Loading model…[/]", end="  ")

    llm = LlamaClass(
        model_path=str(model_path),
        n_ctx=n_ctx,
        n_gpu_layers=n_gpu_layers,
        verbose=verbose,
    )

    console.print("[green]✓[/]")
    console.print()

    console.print(
        Panel(
            "[dim]Type your message and press Enter.\n"
            "/exit  — quit    /clear — clear history\n"
            "/system <text>  — set system prompt[/]",
            title="[bold]Chat[/]",
            border_style="dim",
            padding=(0, 2),
        )
    )
    console.print()

    messages: list[dict[str, str]] = []
    if system_prompt:
        messages.append({"role": "system", "content": system_prompt})

    while True:
        try:
            user_input = input("  [You] › ").strip()
        except (EOFError, KeyboardInterrupt):
            console.print("\n\n  [dim]Session ended.[/]\n")
            break

        if not user_input:
            continue

        # -- built-in slash commands --
        if user_input == "/exit":
            console.print("\n  [dim]Goodbye.[/]\n")
            break
        if user_input == "/clear":
            messages = [m for m in messages if m["role"] == "system"]
            console.print("  [dim]History cleared.[/]\n")
            continue
        if user_input.startswith("/system "):
            new_sys = user_input[len("/system "):].strip()
            messages = [m for m in messages if m["role"] != "system"]
            messages.insert(0, {"role": "system", "content": new_sys})
            console.print("  [dim]System prompt updated.[/]\n")
            continue

        messages.append({"role": "user", "content": user_input})

        # -- stream the response --
        print("  [AI]  ", end="", flush=True)
        full_response = ""

        try:
            stream = llm.create_chat_completion(
                messages=messages,
                stream=True,
                temperature=0.7,
                max_tokens=max_tokens,
            )
            for chunk in stream:
                delta = chunk["choices"][0]["delta"]
                token = delta.get("content", "")
                if token:
                    print(token, end="", flush=True)
                    full_response += token
        except Exception as exc:  # pylint: disable=broad-except
            console.print(f"\n  [red]Error:[/] {exc}")
            messages.pop()  # remove failed user message
            continue

        print("\n")
        messages.append({"role": "assistant", "content": full_response})


def run_gguf_once(
    model_path: Path,
    prompt: str,
    *,
    n_ctx: int = 128000,
    n_gpu_layers: int = -1,
    max_tokens: int = 128000,
    verbose: bool = False,
) -> str:
    """Run a single inference and return the response text."""
    LlamaClass = load_llama_cpp()  # pylint: disable=invalid-name

    llm = LlamaClass(
        model_path=str(model_path),
        n_ctx=n_ctx,
        n_gpu_layers=n_gpu_layers,
        verbose=verbose,
    )
    result = llm.create_chat_completion(
        messages=[{"role": "user", "content": prompt}],
        max_tokens=max_tokens,
    )
    return result["choices"][0]["message"]["content"]


# ── Transformers ─────────────────────────────────────────────────────────────

def _patch_transformers_generic() -> None:
    """
    Inject compatibility shims into transformers.utils.generic for models
    that use trust_remote_code and were written against an older transformers API.

    In transformers 5.x, OutputRecorder moved to transformers.utils.output_capturing
    and check_model_inputs was removed.  Models like MiniMax-M2.5 still import
    both from transformers.utils.generic, causing an ImportError at load time.
    """
    try:
        import transformers.utils.generic as _generic  # type: ignore
    except ImportError:
        return

    if not hasattr(_generic, "OutputRecorder"):
        try:
            from transformers.utils.output_capturing import OutputRecorder  # type: ignore
        except ImportError:
            try:
                from transformers.modeling_utils import OutputRecorder  # type: ignore
            except ImportError:
                # Provide a minimal stub so the attribute assignment in model code
                # doesn't crash.  Actual output recording won't work, but inference will.
                class OutputRecorder:  # type: ignore[no-redef]
                    def __init__(self, module, index=0, layer_name=None):
                        self.module = module
                        self.index = index
                        self.layer_name = layer_name
        _generic.OutputRecorder = OutputRecorder

    if not hasattr(_generic, "check_model_inputs"):
        # check_model_inputs was removed in transformers 5.x.
        # It was a decorator that validated forward() arguments; a no-op
        # passthrough is safe for inference.
        def check_model_inputs(fn):  # type: ignore[no-redef]
            return fn
        _generic.check_model_inputs = check_model_inputs


def load_transformers():
    """Return the transformers module, or exit with an install hint."""
    try:
        # pylint: disable-next=import-outside-toplevel
        import transformers  # type: ignore
        _patch_transformers_generic()
        return transformers
    except ImportError:
        console.print(
            "\n  [bold red]transformers[/] is not installed.\n\n"
            "  Install it with:\n\n"
            "    [bold cyan]pip install transformers torch accelerate[/]\n"
        )
        sys.exit(1)


def _extract_response(generated_text) -> str:
    """Pull the assistant reply out of pipeline output."""
    if isinstance(generated_text, list):
        return generated_text[-1].get("content", "")
    return generated_text


def build_pipe_input(
    pipe,  # type: ignore[no-untyped-def]
    messages: list[dict[str, str]],
) -> "list[dict[str, str]] | str":
    """
    Return the messages list for models that have a chat template,
    or a plain-text fallback for models that don't.
    """
    tokenizer = getattr(pipe, "tokenizer", None)
    if tokenizer and getattr(tokenizer, "chat_template", None):
        return messages
    lines = [f"{m['role'].capitalize()}: {m['content']}" for m in messages]
    lines.append("Assistant:")
    return "\n".join(lines)


def _check_fp8_gpu_requirement(model_dir: Path) -> None:
    """Exit early with a clear message if model needs FP8 GPU but none is available."""
    import json  # pylint: disable=import-outside-toplevel
    config_path = model_dir / "config.json"
    if not config_path.exists():
        return
    try:
        cfg = json.loads(config_path.read_text(encoding="utf-8"))
    except Exception:  # pylint: disable=broad-except
        return
    if cfg.get("quantization_config", {}).get("quant_method") != "fp8":
        return
    try:
        import torch  # type: ignore  # pylint: disable=import-outside-toplevel
        has_gpu = torch.cuda.is_available() or (
            hasattr(torch, "xpu") and torch.xpu.is_available()
        )
    except ImportError:
        has_gpu = False
    if not has_gpu:
        console.print(
            f"\n  [red]error[/]  [bold]{model_dir.name}[/] uses FP8 quantization "
            "which requires a CUDA or XPU GPU.\n\n"
            "  No GPU was detected.  Options:\n"
            "    • Run on a machine with a compatible GPU\n"
            "    • Use a non-quantized or GGUF variant of this model\n"
        )
        sys.exit(1)


def run_transformers_interactive(
    model_dir: Path,
    *,
    device: str = "auto",
    max_tokens: int = 128000,
    system_prompt: str | None = None,
) -> None:
    """Start an interactive chat session with a Transformers model."""
    transformers = load_transformers()
    _check_fp8_gpu_requirement(model_dir)
    from rich.panel import Panel  # pylint: disable=import-outside-toplevel

    console.print("\n  [dim]Loading model…[/]", end="  ")
    try:
        pipe = transformers.pipeline(
            "text-generation",
            model=str(model_dir),
            device_map=device,
            trust_remote_code=True,
        )
    except ValueError as exc:
        msg = str(exc)
        if "Unrecognized model" in msg or "model_type" in msg:
            console.print(
                f"\n  [red]error[/]  [bold]{model_dir.name}[/] is not a "
                "text-generation model.\n\n"
                "  [dim]llmpm run only supports causal language models "
                "(e.g. Llama, Mistral, Qwen, SmolLM…).[/]\n"
            )
        else:
            console.print(f"\n  [red]Failed to load model:[/] {exc}")
        sys.exit(1)
    except Exception as exc:  # pylint: disable=broad-except
        console.print(f"\n  [red]Failed to load model:[/] {exc}")
        sys.exit(1)

    if pipe.tokenizer.pad_token_id is None:
        pipe.tokenizer.pad_token_id = pipe.tokenizer.eos_token_id

    console.print("[green]✓[/]")
    console.print()

    console.print(
        Panel(
            "[dim]Type your message and press Enter.\n"
            "/exit  — quit    /clear — clear history\n"
            "/system <text>  — set system prompt[/]",
            title="[bold]Chat[/]",
            border_style="dim",
            padding=(0, 2),
        )
    )
    console.print()

    history: list[dict[str, str]] = []
    if system_prompt:
        history.append({"role": "system", "content": system_prompt})

    while True:
        try:
            user_input = input("  [You] › ").strip()
        except (EOFError, KeyboardInterrupt):
            console.print("\n\n  [dim]Session ended.[/]\n")
            break

        if not user_input:
            continue
        if user_input == "/exit":
            console.print("\n  [dim]Goodbye.[/]\n")
            break
        if user_input == "/clear":
            history = [m for m in history if m["role"] == "system"]
            console.print("  [dim]History cleared.[/]\n")
            continue
        if user_input.startswith("/system "):
            new_sys = user_input[len("/system "):].strip()
            history = [m for m in history if m["role"] != "system"]
            history.insert(0, {"role": "system", "content": new_sys})
            console.print("  [dim]System prompt updated.[/]\n")
            continue

        history.append({"role": "user", "content": user_input})

        try:
            output = pipe(
                build_pipe_input(pipe, history),
                max_new_tokens=max_tokens,
                return_full_text=False,
                do_sample=True,
                temperature=0.7,
            )
            response = _extract_response(output[0]["generated_text"])
        except Exception as exc:  # pylint: disable=broad-except
            console.print(f"  [red]Error:[/] {exc}\n")
            history.pop()
            continue

        print(f"  [AI]  {response}\n")
        history.append({"role": "assistant", "content": response})


def run_transformers_once(
    model_dir: Path,
    prompt: str,
    *,
    device: str = "auto",
    max_new_tokens: int = 128000,
) -> str:
    """Run a single inference and return the response text."""
    transformers = load_transformers()
    pipe = transformers.pipeline(
        "text-generation",
        model=str(model_dir),
        device_map=device,
        trust_remote_code=True,
    )
    messages = [{"role": "user", "content": prompt}]
    output = pipe(
        build_pipe_input(pipe, messages),
        max_new_tokens=max_new_tokens,
        return_full_text=False,
    )
    return _extract_response(output[0]["generated_text"])