from __future__ import annotations

import platform
import shlex
import shutil
import subprocess
import sys
from dataclasses import dataclass

import typer
from rich.console import Console

SCISPACY_MODEL_URL = (
    "https://s3-us-west-2.amazonaws.com/ai2-s2-scispacy/releases/v0.5.4/en_core_sci_lg-0.5.4.tar.gz"
)


@dataclass(frozen=True)
class SetupStep:
    """One installation step. argv=None means the user must do it manually."""

    label: str
    argv: tuple[str, ...] | None = None
    note: str | None = None


def _tesseract_steps() -> list[SetupStep]:
    if shutil.which("tesseract"):
        return []
    system = platform.system().lower()
    if system == "darwin":
        if shutil.which("brew"):
            return [SetupStep("Install Tesseract via Homebrew", ("brew", "install", "tesseract"))]
        return [
            SetupStep(
                "Install Tesseract",
                None,
                "Install Homebrew first, then run: brew install tesseract",
            )
        ]
    if system == "windows":
        if shutil.which("choco"):
            return [
                SetupStep("Install Tesseract via Chocolatey", ("choco", "install", "tesseract"))
            ]
        return [
            SetupStep(
                "Install Tesseract",
                None,
                "Download the installer from https://github.com/UB-Mannheim/tesseract/wiki",
            )
        ]
    if shutil.which("apt-get"):
        return [
            SetupStep(
                "Install Tesseract via apt", ("sudo", "apt-get", "install", "-y", "tesseract-ocr")
            )
        ]
    return [
        SetupStep(
            "Install Tesseract", None, "Install the tesseract-ocr package for your distribution."
        )
    ]


def steps_for(target: str) -> list[SetupStep]:
    if target == "basic":
        steps = _tesseract_steps()
        steps.append(
            SetupStep("Verify the install", None, "Run: ocr-pipeline init && ocr-pipeline demo")
        )
        return steps
    if target == "full":
        steps = []
        if shutil.which("uv"):
            steps.append(
                SetupStep("Install optional ML dependencies", ("uv", "sync", "--extra", "full"))
            )
            steps.append(
                SetupStep(
                    "Download the scispaCy model (en_core_sci_lg)",
                    ("uv", "pip", "install", SCISPACY_MODEL_URL),
                    "Requires a scispaCy-compatible spaCy; see the README if this step fails.",
                )
            )
        else:
            steps.append(
                SetupStep(
                    "Install optional ML dependencies",
                    (sys.executable, "-m", "pip", "install", "-e", ".[full]"),
                )
            )
            steps.append(
                SetupStep(
                    "Download the scispaCy model (en_core_sci_lg)",
                    (sys.executable, "-m", "pip", "install", SCISPACY_MODEL_URL),
                    "Requires a scispaCy-compatible spaCy; see the README if this step fails.",
                )
            )
        steps.append(
            SetupStep(
                "Detectron2 for figure detection",
                None,
                "No universal wheel exists; install a build matching your torch/Python, see README.",
            )
        )
        return steps
    raise ValueError(f"Unknown setup target: {target}. Choose basic or full.")


def run_steps(steps: list[SetupStep], *, assume_yes: bool, dry_run: bool, console: Console) -> int:
    """Print and optionally execute each step. Returns the first failing exit code."""
    for index, step in enumerate(steps, 1):
        console.print(f"[cyan][{index}/{len(steps)}] {step.label}[/cyan]")
        if step.argv is None:
            if step.note:
                console.print(f"    {step.note}")
            continue
        console.print(f"    $ {shlex.join(step.argv)}")
        if dry_run:
            continue
        if step.note:
            console.print(f"    {step.note}")
        if not assume_yes and not typer.confirm("Run this step?", default=True):
            console.print("    [dim]skipped[/dim]")
            continue
        returncode = subprocess.run(list(step.argv)).returncode
        if returncode != 0:
            console.print(f"[red]Step failed with exit code {returncode}.[/red]")
            return returncode
    return 0
