"""
llmpm run <repo_id>

Runs a locally installed model in an interactive session (or single
inference with --prompt).

Examples:
  llmpm run bartowski/Llama-3.2-1B-Instruct-GGUF
  llmpm run bartowski/Llama-3.2-1B-Instruct-GGUF --prompt "Hello!"
  llmpm run meta-llama/Llama-3.2-1B-Instruct --system "You are a pirate."
  llmpm run amused/amused-256 --prompt "summer in the mountains"
  llmpm run openai/whisper-base --prompt audio.wav
"""

from __future__ import annotations

import sys
import time
from pathlib import Path

import click
from rich.console import Console
from rich.panel import Panel

from llmpm import display
from llmpm.commands._picker import pick_installed_model
from llmpm.core import model_detector, registry, runner
from llmpm.core.downloader import normalize_repo_id

console = Console(stderr=True)


@click.command("run")
@click.argument("repo_id", required=False, default=None)
@click.option(
    "--prompt", "-p",
    default=None,
    help="Single inference prompt (non-interactive).",
)
@click.option(
    "--system", "-s",
    default=None,
    help="System prompt for the conversation (text models only).",
)
@click.option(
    "--ctx", "n_ctx",
    default=4096,
    show_default=True,
    help="Context window size (GGUF only).",
)
@click.option(
    "--gpu-layers", "n_gpu_layers",
    default=-1,
    show_default=True,
    help="Number of layers to offload to GPU. -1 = auto (GGUF only).",
)
@click.option(
    "--verbose", is_flag=True,
    help="Show verbose model loading output.",
)
@click.option(
    "--max-tokens", "max_tokens",
    default=2048,
    show_default=True,
    help="Maximum tokens to generate per response.",
)
@click.option(
    "--path",
    "model_path_override",
    default=None,
    type=click.Path(exists=True, file_okay=True, dir_okay=True),
    help="Path to a local model directory or GGUF file (bypasses registry).",
)
def run_model(  # pylint: disable=too-many-arguments
    repo_id: str | None,
    prompt: str | None,
    system: str | None,
    n_ctx: int,
    n_gpu_layers: int,
    verbose: bool,
    max_tokens: int,
    model_path_override: str | None,
) -> None:
    """Run an installed model interactively or with a single prompt."""
    if repo_id:
        repo_id = normalize_repo_id(repo_id)

    if model_path_override:
        model_path, model_type, primary_file = _resolve_path_override(
            model_path_override
        )
        repo_id = repo_id or model_path.name
        display.info(
            f"Running model from local path: [bold]{model_path}[/]"
        )
    else:
        if not repo_id:
            display.error(
                "Provide a model name or use "
                "[bold cyan]--path /path/to/model[/]."
            )
            raise SystemExit(1)

        matches = registry.find_models(repo_id)
        if not matches:
            display.error(
                f"No installed model matches [bold]{repo_id}[/].\n\n"
                f"  Run:  [bold cyan]llmpm install {repo_id}[/]"
            )
            raise SystemExit(1)

        if len(matches) == 1:
            matched_id, entry = matches[0]
            if matched_id != repo_id:
                display.info(f"Using [bold]{matched_id}[/]")
        else:
            result = pick_installed_model(matches, repo_id)
            if result is None:
                raise SystemExit(0)
            matched_id, entry = result

        repo_id = matched_id
        model_type = entry["model_type"]
        model_path = Path(entry["path"])
        primary_file = entry.get("primary_file")

        if not model_path.exists():
            display.error(
                f"Model [bold]{repo_id}[/] is registered but its files are missing "
                f"at [bold]{model_path}[/].\n\n"
                f"  Re-install with:  [bold cyan]llmpm install {repo_id}[/]"
            )
            raise SystemExit(1)

    display.header(f"run {repo_id}")
    display.blank()

    display.run_header(
        repo_id=repo_id,
        model_type=model_type,
        primary_file=primary_file,
        n_ctx=n_ctx if model_type == "gguf" else None,
        n_gpu_layers=n_gpu_layers if model_type == "gguf" else None,
        max_tokens=max_tokens,
        system_prompt=system,
    )

    # ── GGUF ──────────────────────────────────────────────────────────────────
    if model_type == "gguf":
        gguf_path = _resolve_gguf_path(model_path, primary_file)
        if gguf_path is None:
            display.error(
                "Cannot find the GGUF file. "
                f"Try re-installing: [bold cyan]llmpm install {repo_id}[/]"
            )
            raise SystemExit(1)

        if prompt:
            result = runner.run_gguf_once(
                gguf_path, prompt,
                n_ctx=n_ctx, n_gpu_layers=n_gpu_layers,
                max_tokens=max_tokens, verbose=verbose,
            )
            print(result)
        else:
            runner.run_gguf_interactive(
                gguf_path,
                n_ctx=n_ctx, n_gpu_layers=n_gpu_layers,
                max_tokens=max_tokens, verbose=verbose,
                system_prompt=system,
            )
        return

    # ── HuggingFace: detect pipeline category ─────────────────────────────────
    category = model_detector.detect(model_path)
    display.info(f"Detected: [bold]{model_detector.label(category)}[/]")

    if category == model_detector.IMAGE_GENERATION:
        if prompt:
            _run_diffusion_once(model_path, prompt)
        else:
            _run_diffusion_interactive(model_path)

    elif category == model_detector.VISION:
        if prompt:
            _run_vision_once(model_path, prompt)
        else:
            _run_vision_interactive(model_path)

    elif category == model_detector.AUDIO_TRANSCRIPTION:
        audio_path = prompt  # --prompt is the audio file path for ASR
        if audio_path:
            _run_asr_once(model_path, audio_path)
        else:
            _run_asr_interactive(model_path)

    elif category == model_detector.AUDIO_SPEECH:
        if prompt:
            _run_tts_once(model_path, prompt)
        else:
            _run_tts_interactive(model_path)

    else:
        if prompt:
            result = runner.run_transformers_once(
                model_path, prompt, max_new_tokens=max_tokens,
            )
            print(result)
        else:
            runner.run_transformers_interactive(
                model_path, max_tokens=max_tokens, system_prompt=system,
            )


# ── image generation ──────────────────────────────────────────────────────────

def _load_diffusion(model_path: Path):
    """Load a diffusion pipeline, returns (pipe, device)."""
    try:
        import torch  # type: ignore  # pylint: disable=import-outside-toplevel
        from diffusers import (  # type: ignore  # pylint: disable=import-outside-toplevel
            AutoPipelineForText2Image, DiffusionPipeline,
        )
    except ImportError as exc:
        console.print(f"\n  [red]Missing dependency:[/] {exc}\n"
                      "  Run: [bold cyan]pip install diffusers torch[/]\n")
        sys.exit(1)

    if torch.cuda.is_available():
        device, dtype = "cuda", torch.float16
    elif getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
        device, dtype = "mps", torch.float16
    else:
        device, dtype = "cpu", torch.float32

    console.print("\n  [dim]Loading model…[/]", end="  ")
    try:
        try:
            pipe = AutoPipelineForText2Image.from_pretrained(
                str(model_path), torch_dtype=dtype,
            )
        except ValueError:
            pipe = DiffusionPipeline.from_pretrained(
                str(model_path), torch_dtype=dtype,
            )
        pipe = pipe.to(device)
    except Exception as exc:  # pylint: disable=broad-except
        console.print(f"\n  [red]Failed to load model:[/] {exc}\n")
        sys.exit(1)

    console.print("[green]✓[/]\n")
    return pipe


def _save_image(b64: str) -> Path:
    """Decode base64 PNG and save to ~/Desktop (fallback: cwd)."""
    import base64  # pylint: disable=import-outside-toplevel
    desktop = Path.home() / "Desktop"
    out_dir = desktop if desktop.exists() else Path.cwd()
    out_path = out_dir / f"llmpm_{int(time.time())}.png"
    out_path.write_bytes(base64.b64decode(b64))
    return out_path


def _run_diffusion_once(model_path: Path, prompt: str) -> None:
    import base64, io  # pylint: disable=import-outside-toplevel,multiple-imports
    pipe = _load_diffusion(model_path)
    console.print("  [dim]Generating…[/]")
    result = pipe(prompt=prompt)
    buf = io.BytesIO()
    result.images[0].save(buf, format="PNG")
    out = _save_image(base64.b64encode(buf.getvalue()).decode())
    console.print(f"  [green]✓[/]  Saved → [bold]{out}[/]\n")


def _run_diffusion_interactive(model_path: Path) -> None:
    import base64, io  # pylint: disable=import-outside-toplevel,multiple-imports
    pipe = _load_diffusion(model_path)
    console.print(
        Panel(
            "[dim]Type an image description and press Enter.\n"
            "/exit — quit[/]",
            title="[bold]Image Generation[/]",
            border_style="dim", padding=(0, 2),
        )
    )
    console.print()
    while True:
        try:
            prompt = input("  [Prompt] › ").strip()
        except (EOFError, KeyboardInterrupt):
            console.print("\n\n  [dim]Session ended.[/]\n")
            break
        if not prompt:
            continue
        if prompt == "/exit":
            console.print("\n  [dim]Goodbye.[/]\n")
            break
        console.print("  [dim]Generating…[/]", end="  ")
        result = pipe(prompt=prompt)
        buf = io.BytesIO()
        result.images[0].save(buf, format="PNG")
        out = _save_image(base64.b64encode(buf.getvalue()).decode())
        console.print(f"[green]✓[/]  Saved → [bold]{out}[/]\n")


# ── vision ────────────────────────────────────────────────────────────────────

def _load_vision(model_path: Path):
    """Load an image-to-text pipeline."""
    try:
        from transformers import (  # type: ignore  # pylint: disable=import-outside-toplevel
            pipeline as hf_pipeline,
        )
    except ImportError:
        console.print("\n  [red]transformers[/] is not installed.\n"
                      "  Run: [bold cyan]pip install transformers torch Pillow[/]\n")
        sys.exit(1)
    console.print("\n  [dim]Loading model…[/]", end="  ")
    try:
        pipe = hf_pipeline(
            "image-to-text", model=str(model_path),
            device_map="auto", trust_remote_code=True,
        )
    except Exception as exc:  # pylint: disable=broad-except
        console.print(f"\n  [red]Failed to load model:[/] {exc}\n")
        sys.exit(1)
    console.print("[green]✓[/]\n")
    return pipe


def _run_vision_once(model_path: Path, image_path: str) -> None:
    pipe = _load_vision(model_path)
    result = pipe(image_path)
    text = result[0].get("generated_text", "") if result else ""
    print(text)


def _run_vision_interactive(model_path: Path) -> None:
    pipe = _load_vision(model_path)
    console.print(
        Panel(
            "[dim]Enter an image file path or URL, then press Enter.\n"
            "/exit — quit[/]",
            title="[bold]Vision (Image → Text)[/]",
            border_style="dim", padding=(0, 2),
        )
    )
    console.print()
    while True:
        try:
            img_input = input("  [Image path/URL] › ").strip()
        except (EOFError, KeyboardInterrupt):
            console.print("\n\n  [dim]Session ended.[/]\n")
            break
        if not img_input:
            continue
        if img_input == "/exit":
            console.print("\n  [dim]Goodbye.[/]\n")
            break
        try:
            result = pipe(img_input)
            text = result[0].get("generated_text", "") if result else "(no output)"
            console.print(f"  [AI]  {text}\n")
        except Exception as exc:  # pylint: disable=broad-except
            console.print(f"  [red]Error:[/] {exc}\n")


# ── audio transcription ───────────────────────────────────────────────────────

def _load_asr(model_path: Path):
    """Load an ASR pipeline."""
    try:
        from transformers import (  # type: ignore  # pylint: disable=import-outside-toplevel
            pipeline as hf_pipeline,
        )
    except ImportError:
        console.print("\n  [red]transformers[/] is not installed.\n"
                      "  Run: [bold cyan]pip install transformers torch[/]\n")
        sys.exit(1)
    console.print("\n  [dim]Loading model…[/]", end="  ")
    try:
        pipe = hf_pipeline(
            "automatic-speech-recognition", model=str(model_path),
            device_map="auto", trust_remote_code=True,
        )
    except Exception as exc:  # pylint: disable=broad-except
        console.print(f"\n  [red]Failed to load model:[/] {exc}\n")
        sys.exit(1)
    console.print("[green]✓[/]\n")
    return pipe


def _run_asr_once(model_path: Path, audio_path: str) -> None:
    pipe = _load_asr(model_path)
    result = pipe(audio_path)
    print(result.get("text", ""))


def _run_asr_interactive(model_path: Path) -> None:
    pipe = _load_asr(model_path)
    console.print(
        Panel(
            "[dim]Enter an audio file path and press Enter to transcribe.\n"
            "/exit — quit[/]",
            title="[bold]Audio Transcription (ASR)[/]",
            border_style="dim", padding=(0, 2),
        )
    )
    console.print()
    while True:
        try:
            audio_path = input("  [Audio file] › ").strip()
        except (EOFError, KeyboardInterrupt):
            console.print("\n\n  [dim]Session ended.[/]\n")
            break
        if not audio_path:
            continue
        if audio_path == "/exit":
            console.print("\n  [dim]Goodbye.[/]\n")
            break
        try:
            result = pipe(audio_path)
            console.print(f"  [AI]  {result.get('text', '(empty)')}\n")
        except Exception as exc:  # pylint: disable=broad-except
            console.print(f"  [red]Error:[/] {exc}\n")


# ── text-to-speech ────────────────────────────────────────────────────────────

def _load_tts(model_path: Path):
    """Load a TTS pipeline."""
    try:
        from transformers import (  # type: ignore  # pylint: disable=import-outside-toplevel
            pipeline as hf_pipeline,
        )
    except ImportError:
        console.print("\n  [red]transformers[/] is not installed.\n"
                      "  Run: [bold cyan]pip install transformers torch numpy[/]\n")
        sys.exit(1)
    console.print("\n  [dim]Loading model…[/]", end="  ")
    try:
        pipe = hf_pipeline(
            "text-to-speech", model=str(model_path),
            device_map="auto", trust_remote_code=True,
        )
    except Exception as exc:  # pylint: disable=broad-except
        console.print(f"\n  [red]Failed to load model:[/] {exc}\n")
        sys.exit(1)
    console.print("[green]✓[/]\n")
    return pipe


def _save_wav(pipe_result) -> Path:
    """Convert TTS output to WAV and save to Desktop or cwd."""
    import io, struct, wave  # pylint: disable=import-outside-toplevel,multiple-imports
    import numpy as np  # type: ignore  # pylint: disable=import-outside-toplevel
    audio = pipe_result["audio"]
    if hasattr(audio, "numpy"):
        audio = audio.numpy()
    audio_arr = np.array(audio, dtype=np.float32).flatten()
    sr = pipe_result.get("sampling_rate", 22050)
    int16 = (audio_arr * 32767).astype(np.int16)
    buf = io.BytesIO()
    with wave.open(buf, "wb") as wf:
        wf.setnchannels(1)
        wf.setsampwidth(2)
        wf.setframerate(sr)
        wf.writeframes(struct.pack(f"<{len(int16)}h", *int16))
    desktop = Path.home() / "Desktop"
    out_dir = desktop if desktop.exists() else Path.cwd()
    out_path = out_dir / f"llmpm_{int(time.time())}.wav"
    out_path.write_bytes(buf.getvalue())
    return out_path


def _run_tts_once(model_path: Path, text: str) -> None:
    pipe = _load_tts(model_path)
    result = pipe(text)
    out = _save_wav(result)
    console.print(f"  [green]✓[/]  Saved → [bold]{out}[/]\n")


def _run_tts_interactive(model_path: Path) -> None:
    pipe = _load_tts(model_path)
    console.print(
        Panel(
            "[dim]Type text to synthesize and press Enter.\n"
            "/exit — quit[/]",
            title="[bold]Text-to-Speech (TTS)[/]",
            border_style="dim", padding=(0, 2),
        )
    )
    console.print()
    while True:
        try:
            text = input("  [Text] › ").strip()
        except (EOFError, KeyboardInterrupt):
            console.print("\n\n  [dim]Session ended.[/]\n")
            break
        if not text:
            continue
        if text == "/exit":
            console.print("\n  [dim]Goodbye.[/]\n")
            break
        try:
            result = pipe(text)
            out = _save_wav(result)
            console.print(f"  [green]✓[/]  Saved → [bold]{out}[/]\n")
        except Exception as exc:  # pylint: disable=broad-except
            console.print(f"  [red]Error:[/] {exc}\n")


# ── helpers ───────────────────────────────────────────────────────────────────

def _resolve_path_override(raw: str) -> tuple[Path, str, str | None]:
    """
    Given a --path value, return (model_dir, model_type, primary_file).

    Accepts:
      - a single .gguf file  → model_type="gguf", primary_file=<filename>
      - a directory with .gguf files inside → model_type="gguf", primary_file=None
      - any other directory → model_type="transformers", primary_file=None
    """
    p = Path(raw).expanduser().resolve()
    if p.is_file():
        if p.suffix.lower() == ".gguf":
            return p.parent, "gguf", p.name
        display.error(f"[bold]{raw}[/] is not a directory or .gguf file.")
        raise SystemExit(1)
    if not p.is_dir():
        display.error(f"[bold]{raw}[/] does not exist.")
        raise SystemExit(1)
    gguf_files = list(p.glob("*.gguf"))
    if gguf_files:
        return p, "gguf", None
    return p, "transformers", None


def _resolve_gguf_path(model_dir: Path, primary_file: str | None) -> Path | None:
    """Find the GGUF file inside the model directory."""
    if primary_file:
        candidate = model_dir / primary_file
        if candidate.exists():
            return candidate
    for gguf_file in model_dir.rglob("*.gguf"):
        return gguf_file
    return None
