from __future__ import annotations

import json
import sys
from dataclasses import dataclass
from typing import Annotated

import typer
from pydantic import ValidationError
from rich import print as rprint

from atlas_runner.errors import ErrorCode, RunnerError
from atlas_runner.models import Mount, RunSpec
from atlas_runner.runner import Runner, create_backend

app = typer.Typer(name="atlas-runner", add_completion=False)

# ------------------------------------------------------------------
# Global state populated by the callback
# ------------------------------------------------------------------


@dataclass
class _State:
    backend: str = "docker"
    json: bool = False
    docker_host: str | None = None


_state = _State()


def _get_runner() -> Runner:
    backend = create_backend(
        backend_type=_state.backend,
        docker_host=_state.docker_host,
    )
    return Runner(backend)


def _json_enabled() -> bool:
    return _state.json


# ------------------------------------------------------------------
# Helpers
# ------------------------------------------------------------------


def _parse_key_value_list(items: list[str]) -> dict[str, str]:
    """Parse a list of ``KEY=VALUE`` strings into a dict."""
    result: dict[str, str] = {}
    for item in items:
        if "=" not in item:
            raise typer.BadParameter(f"Expected KEY=VALUE format, got: {item}")
        key, val = item.split("=", 1)
        result[key] = val
    return result


def _output(data: object) -> None:
    """Print data as JSON or human-readable."""
    if _json_enabled():
        typer.echo(json.dumps(data, default=str))
    else:
        rprint(data)


def _handle_error(err: RunnerError) -> None:
    """Print a RunnerError and exit."""
    if _json_enabled():
        typer.echo(json.dumps(err.to_dict(), default=str))
    else:
        rprint(f"[red]Error[/red] [{err.code.value}]: {err.message}")
        if err.details:
            rprint(err.details)
    raise typer.Exit(code=1)


# ------------------------------------------------------------------
# Global options
# ------------------------------------------------------------------


@app.callback()
def main(
    backend: Annotated[str, typer.Option("--backend", help="Backend type.")] = "docker",
    json_output: Annotated[
        bool, typer.Option("--json", help="Output as JSON.")
    ] = False,
    docker_host: Annotated[
        str | None,
        typer.Option("--docker-host", help="Docker daemon socket."),
    ] = None,
) -> None:
    """Atlas Runner — container orchestration CLI."""
    _state.backend = backend
    _state.json = json_output
    _state.docker_host = docker_host


# ------------------------------------------------------------------
# Commands
# ------------------------------------------------------------------


@app.command()
def login(
    registry: Annotated[str, typer.Argument(help="Registry URL.")],
    username: Annotated[str, typer.Option("--username", "-u", help="Username.")],
    password: Annotated[
        str | None,
        typer.Option("--password", "-p", help="Password."),
    ] = None,
    password_stdin: Annotated[
        bool,
        typer.Option("--password-stdin", help="Read password from stdin."),
    ] = False,
) -> None:
    """Authenticate with a container registry."""
    if password is not None and password_stdin:
        raise typer.BadParameter(
            "Specify either --password or --password-stdin, not both."
        )
    if password_stdin:
        password = sys.stdin.read().strip()
    if not password:
        _handle_error(
            RunnerError(
                code=ErrorCode.LOGIN_FAILED,
                message="password required (use --password or --password-stdin)",
            )
        )
        return
    try:
        runner = _get_runner()
        runner.login(registry, username, password)
        _output({"status": "ok", "registry": registry})
    except RunnerError as e:
        _handle_error(e)


@app.command()
def pull(
    image: Annotated[str, typer.Argument(help="Image to pull.")],
) -> None:
    """Pull a container image."""
    try:
        runner = _get_runner()
        runner.pull(image)
        _output({"status": "ok", "image": image})
    except RunnerError as e:
        _handle_error(e)


@app.command()
def run(
    name: Annotated[str, typer.Option("--name", help="Container name.")],
    image: Annotated[str, typer.Option("--image", help="Container image.")],
    cmd: Annotated[
        list[str] | None,
        typer.Option("--cmd", help="Command arguments (repeatable)."),
    ] = None,
    shell: Annotated[
        str | None,
        typer.Option("--shell", help="Shell command string."),
    ] = None,
    env: Annotated[
        list[str] | None,
        typer.Option("--env", "-e", help="Environment variable KEY=VALUE."),
    ] = None,
    label: Annotated[
        list[str] | None,
        typer.Option("--label", "-l", help="Label KEY=VALUE."),
    ] = None,
    gpu: Annotated[
        list[int] | None,
        typer.Option("--gpu", help="GPU device index (repeatable)."),
    ] = None,
    cpu_cores: Annotated[
        int | None,
        typer.Option("--cpu-cores", help="Number of CPU cores."),
    ] = None,
    cpuset: Annotated[
        str | None,
        typer.Option("--cpuset", help="CPU set (e.g. 0-3)."),
    ] = None,
    memory_gb: Annotated[
        int | None,
        typer.Option("--memory-gb", help="Memory limit in GB."),
    ] = None,
    scratch: Annotated[
        str | None,
        typer.Option("--scratch", help="Host path for scratch volume."),
    ] = None,
    mount: Annotated[
        list[str] | None,
        typer.Option("--mount", help="Mount SRC:DST[:ro]."),
    ] = None,
    network: Annotated[
        str | None,
        typer.Option("--network", help="Docker network."),
    ] = None,
    port: Annotated[
        list[str] | None,
        typer.Option("--port", help="Port mapping HOST:CONTAINER."),
    ] = None,
    workdir: Annotated[
        str | None,
        typer.Option("--workdir", help="Working directory."),
    ] = None,
) -> None:
    """Run a container."""
    try:
        mounts: list[Mount] = []
        for m in mount or []:
            parts = m.split(":")
            if len(parts) not in (2, 3):
                raise typer.BadParameter(
                    f"Invalid mount format: {m}. Expected SRC:DST[:ro]."
                )
            if len(parts) == 3 and parts[2] != "ro":
                raise typer.BadParameter(
                    f"Unsupported mount flag in {m!r}: only 'ro' is allowed."
                )
            ro = len(parts) == 3
            mounts.append(Mount(src=parts[0], dst=parts[1], ro=ro))

        ports: dict[str, int] = {}
        for p in port or []:
            if ":" not in p:
                raise typer.BadParameter(f"Invalid port format: {p}")
            host_p, container_p = p.split(":", 1)
            try:
                ports[host_p] = int(container_p)
            except ValueError:
                raise typer.BadParameter(
                    f"Invalid port format: {p}. Container port must be an integer."
                )

        try:
            spec = RunSpec(
                name=name,
                image=image,
                command=cmd if cmd else None,
                shell=shell,
                labels=_parse_key_value_list(label or []),
                gpus=gpu if gpu else None,
                cpu_cores=cpu_cores,
                cpuset=cpuset,
                memory_gb=memory_gb,
                scratch=scratch,
                mounts=mounts,
                network=network,
                ports=ports,
                workdir=workdir,
                env=_parse_key_value_list(env or []),
            )
        except ValidationError as e:
            raise typer.BadParameter(str(e)) from e
        runner = _get_runner()
        result = runner.run(spec)
        _output(result.model_dump())
    except RunnerError as e:
        _handle_error(e)


@app.command()
def status(
    name: Annotated[str, typer.Argument(help="Container name.")],
) -> None:
    """Get container status."""
    try:
        runner = _get_runner()
        result = runner.status(name)
        _output(result.model_dump())
    except RunnerError as e:
        _handle_error(e)


@app.command()
def stop(
    name: Annotated[str, typer.Argument(help="Container name.")],
) -> None:
    """Stop a container."""
    try:
        runner = _get_runner()
        runner.stop(name)
        _output({"status": "ok", "name": name})
    except RunnerError as e:
        _handle_error(e)


@app.command()
def destroy(
    name: Annotated[str, typer.Argument(help="Container name.")],
) -> None:
    """Stop and remove a container."""
    try:
        runner = _get_runner()
        runner.destroy(name)
        _output({"status": "ok", "name": name})
    except RunnerError as e:
        _handle_error(e)


@app.command()
def logs(
    name: Annotated[str, typer.Argument(help="Container name.")],
) -> None:
    """Fetch container logs."""
    try:
        runner = _get_runner()
        output = runner.logs(name)
        if _json_enabled():
            typer.echo(json.dumps({"name": name, "logs": output}))
        else:
            typer.echo(output)
    except RunnerError as e:
        _handle_error(e)


@app.command()
def ps() -> None:
    """List containers."""
    try:
        runner = _get_runner()
        results = runner.ps()
        _output([r.model_dump() for r in results])
    except RunnerError as e:
        _handle_error(e)
