"""Vast.ai instance fetching, naming, display, and selection."""

from __future__ import annotations

import json
import re
import subprocess
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import vastly
from vastly.config import Config
from vastly import dim
from vastly.errors import APIError, VastlyError
from vastly.ssh import (
    SSH_CONFIG_DIR,
    ensure_ssh_include,
    find_available_port,
    prune_ssh_configs,
    write_ssh_config,
)


@dataclass
class Instance:
    """A Vast.ai GPU instance."""

    name: str
    id: int
    dph_total: float
    gpu_name: str
    num_gpus: int
    status: str  # raw API state ("running", "stopped", "exited", etc.)
    alias: str | None

    @property
    def display_name(self) -> str:
        """Return alias (auto-name) if aliased, otherwise just name."""
        return f"{self.alias} ({self.name})" if self.alias else self.name


_ALIASES_FILE = Path.home() / ".vastly" / "aliases.json"

# Reserved names that cannot be used as aliases
_RESERVED_NAMES = frozenset(
    {
        "connect",
        "stop",
        "destroy",
        "cp",
        "name",
        "list",
        "start",
        "config",
        "ssh",
    }
)

# Instance lifecycle state classifications
STOPPED_STATES = {"stopped", "exited"}
TRANSITIONAL_STATES = {"loading", "creating", "connecting"}
STARTABLE_STATES = STOPPED_STATES | TRANSITIONAL_STATES
STOPPABLE_STATES = {"running"} | TRANSITIONAL_STATES


def fetch_instances() -> list[dict[str, Any]]:
    """Call vastai CLI and return raw instance data for all instances.

    Raises APIError if the API is unreachable or returns invalid data.
    """
    try:
        result = subprocess.run(
            ["vastai", "show", "instances", "--raw"],
            capture_output=True,
            text=True,
            timeout=30,
        )
    except subprocess.TimeoutExpired:
        raise APIError("vastai command timed out (network issue or Vast.ai outage).")
    if result.returncode != 0:
        if result.stderr.strip():
            raise APIError(f"vastai: {result.stderr.strip()}")
        raise APIError(
            "vastai command failed (missing API key, network issue, or Vast.ai outage). "
            "Run 'vastai set api-key <key>' if you haven't set one."
        )
    try:
        data = json.loads(result.stdout)
    except (json.JSONDecodeError, TypeError) as e:
        raise APIError(
            "vastai returned invalid data. Try running 'vastai show instances --raw' to debug."
        ) from e

    if not isinstance(data, list):
        raise APIError(
            "vastai returned invalid data. Try running 'vastai show instances --raw' to debug."
        )
    return data


def build_instance_name(inst: dict[str, Any], seen: set[str]) -> str:
    """Generate a human-readable name like '1xrtx4090-tw'.

    Names are lowercased to stay compatible with tools (e.g. VS Code
    Remote-SSH) that normalise SSH host aliases to lowercase.

    Appends the instance ID on collision.
    """
    gpu_name = re.sub(r"[^a-zA-Z0-9]", "", inst.get("gpu_name", "")) or "unknown"
    num_gpus = inst.get("num_gpus", 1)
    geo = inst.get("geolocation", "")

    m = re.search(r",\s*(\w+)$", geo)
    country = m.group(1) if m else ""

    base = f"{num_gpus}x{gpu_name}-{country}" if country else f"{num_gpus}x{gpu_name}"
    base = base.lower()

    if base in seen:
        # Collision: append ID to disambiguate (don't add to seen -- we
        # track base names only so the *first* user keeps the short name).
        name = f"{base}-{inst['id']}"
    else:
        seen.add(base)
        name = base

    return name


def load_aliases() -> dict[str, str]:
    """Load aliases from ~/.vastly/aliases.json.

    Returns a mapping of instance ID (string) to alias name.
    """
    if not _ALIASES_FILE.exists():
        return {}
    try:
        return json.loads(_ALIASES_FILE.read_text(encoding="utf-8"))
    except (json.JSONDecodeError, OSError):
        vastly.verbose(f"Warning: could not read {_ALIASES_FILE}, using empty aliases")
        return {}


def save_aliases(aliases: dict[str, str]) -> None:
    """Save aliases to ~/.vastly/aliases.json."""
    _ALIASES_FILE.parent.mkdir(parents=True, exist_ok=True)
    _ALIASES_FILE.write_text(json.dumps(aliases, indent=2) + "\n", encoding="utf-8")


def validate_alias(
    alias: str, instances: list[Instance], aliases: dict[str, str]
) -> None:
    """Validate an alias name. Raises VastlyError on invalid input."""
    if not alias:
        raise VastlyError("Alias cannot be empty.")
    if not re.match(r"^[a-z0-9][a-z0-9-]*$", alias):
        raise VastlyError(
            f"Invalid alias '{alias}'. Must be lowercase alphanumeric with hyphens."
        )
    if alias in _RESERVED_NAMES:
        raise VastlyError(
            f"'{alias}' is a reserved command name and cannot be used as an alias."
        )

    # Check collision with auto-generated names
    auto_names = {i.name for i in instances}
    if alias in auto_names:
        raise VastlyError(f"'{alias}' conflicts with an auto-generated instance name.")

    # Check collision with other aliases
    for inst_id, existing_alias in aliases.items():
        if existing_alias == alias:
            raise VastlyError(
                f"Alias '{alias}' is already assigned to instance {inst_id}."
            )


def find_by_name(instances: list[Instance], name: str) -> Instance | None:
    """Return the instance matching by name or alias, or None."""
    for i in instances:
        if i.name == name or i.alias == name:
            return i
    return None


def sync_instances(config: Config) -> list[Instance]:
    """Sync instances from the API to SSH configs.

    Returns a list of Instance objects for ALL states (running, stopped,
    exited, etc.). SSH configs are only written for running instances.

    Raises APIError if the API is unreachable.
    """
    ensure_ssh_include()
    SSH_CONFIG_DIR.mkdir(parents=True, exist_ok=True)

    all_api = fetch_instances()

    running = [i for i in all_api if i.get("cur_state") == "running"]
    non_running = [i for i in all_api if i.get("cur_state") != "running"]
    vastly.verbose(f"API returned {len(all_api)} instances, {len(running)} running")

    seen: set[str] = set()
    used_ports: set[int] = set()
    results: list[Instance] = []
    # SSH config params for running instances, keyed by instance ID.
    # Used later to write alias configs without leaking internal data onto Instance.
    ssh_params: dict[int, dict] = {}

    # Process running instances first (they get SSH configs and clean names)
    for inst in running:
        name = build_instance_name(inst, seen)

        # Get SSH port -- if no ports yet, the instance is still starting
        try:
            ssh_port = inst["ports"]["22/tcp"][0]["HostPort"]
        except (KeyError, IndexError, TypeError):
            # cur_state says "running" but no ports -- use actual_status
            actual = inst.get("actual_status", "loading")
            status = actual if actual != "running" else "loading"
            results.append(
                Instance(
                    name=name,
                    id=inst["id"],
                    dph_total=inst.get("dph_total", 0),
                    gpu_name=inst.get("gpu_name", ""),
                    num_gpus=inst.get("num_gpus", 0),
                    status=status,
                    alias=None,
                )
            )
            continue

        local_forwards = []
        for pf in config["portForwards"]:
            local_port = find_available_port(int(pf["local"]), used_ports)
            used_ports.add(local_port)
            local_forwards.append((local_port, int(pf["remote"])))

        write_ssh_config(
            name,
            host=inst["public_ipaddr"],
            port=int(ssh_port),
            user=config["sshUser"],
            key_path=config["sshKeyPath"],
            local_forwards=local_forwards,
        )

        ssh_params[inst["id"]] = {
            "host": inst["public_ipaddr"],
            "port": int(ssh_port),
            "user": config["sshUser"],
            "key_path": config["sshKeyPath"],
            "local_forwards": local_forwards,
        }

        results.append(
            Instance(
                name=name,
                id=inst["id"],
                dph_total=inst.get("dph_total", 0),
                gpu_name=inst.get("gpu_name", ""),
                num_gpus=inst.get("num_gpus", 0),
                status="running",
                alias=None,
            )
        )

    # Process non-running instances (no SSH config, no port extraction)
    for inst in non_running:
        name = build_instance_name(inst, seen)

        results.append(
            Instance(
                name=name,
                id=inst["id"],
                dph_total=inst.get("dph_total", 0),
                gpu_name=inst.get("gpu_name", ""),
                num_gpus=inst.get("num_gpus", 0),
                status=inst.get("cur_state", "unknown"),
                alias=None,
            )
        )

    # Prune aliases for instances that no longer exist in the API at all
    aliases = load_aliases()
    all_api_ids = {str(r.id) for r in results}
    stale_ids = [k for k in aliases if k not in all_api_ids]
    if stale_ids:
        for k in stale_ids:
            del aliases[k]
        save_aliases(aliases)

    # Write alias SSH configs (running only) and attach aliases to all instances
    for r in results:
        inst_id = str(r.id)
        alias = aliases.get(inst_id)
        if alias:
            r.alias = alias
            if r.status == "running" and r.id in ssh_params:
                params = ssh_params[r.id]
                write_ssh_config(
                    alias,
                    host=params["host"],
                    port=params["port"],
                    user=params["user"],
                    key_path=params["key_path"],
                    local_forwards=params["local_forwards"],
                )

    # Prune stale SSH configs (ones we didn't just write)
    written_configs = {
        r.name for r in results if r.status == "running" and r.id in ssh_params
    }
    written_configs |= {
        r.alias
        for r in results
        if r.alias and r.status == "running" and r.id in ssh_params
    }
    prune_ssh_configs(written_configs)

    return results


def get_synced_instances(config: Config) -> list[Instance]:
    """Sync and return all instances (any state).

    Raises VastlyError if no instances exist at all.
    Raises APIError if the API is unreachable.
    """
    result = sync_instances(config)
    if not result:
        raise VastlyError("No Vast instances found.")
    return result


def get_running_instances(config: Config) -> list[Instance]:
    """Sync and return only running instances.

    Raises VastlyError if no running instances exist.
    """
    all_instances = sync_instances(config)
    running = [i for i in all_instances if i.status == "running"]
    if not running:
        if all_instances:
            raise VastlyError(
                f"No running instances. {len(all_instances)} inactive. "
                "Use 'vst list' to see all, 'vst start' to restart."
            )
        raise VastlyError("No Vast instances found.")
    return running


def show_table(instances: list[Instance]) -> None:
    """Print a compact, column-aligned instance table."""
    if not instances:
        print()
        return

    # Build rows: (label, cost, right_col, is_running)
    rows: list[tuple[str, str, str, bool]] = []
    for inst in instances:
        label = inst.display_name
        cost = f"${inst.dph_total:.2f}/hr"
        if inst.status == "running":
            rows.append((label, cost, "running", True))
        else:
            rows.append((label, cost, "inactive", False))

    label_w = max(len(r[0]) for r in rows)
    cost_w = max(len(r[1]) for r in rows)

    for label, cost, right, is_running in rows:
        if is_running:
            print(f"  {label.ljust(label_w)}  {cost.ljust(cost_w)}  {right}")
        else:
            print(dim(f"  {label.ljust(label_w)}  {cost.ljust(cost_w)}  {right}"))
    print()


def select_instance(
    instances: list[Instance],
    name: str | None = None,
    *,
    allow_all: bool = False,
    prompt: str = "Select instance:",
) -> list[Instance]:
    """Select instance(s) by name/alias, auto-select if only one, or prompt.

    When allow_all is True, shows an "[a] All" option for multi-instance selection.
    Always returns a non-empty list. Raises VastlyError on cancel or invalid input.
    """
    if not instances:
        raise VastlyError("No instances available.")

    if name:
        match = find_by_name(instances, name)
        if not match:
            labels = [i.display_name for i in instances]
            raise VastlyError(
                f"No instance named '{name}'. Available: {', '.join(labels)}"
            )
        vastly.verbose(f"Matched instance by name/alias: {match.display_name}")
        return [match]

    if len(instances) == 1:
        vastly.verbose(f"Auto-selected sole instance: {instances[0].display_name}")
        return [instances[0]]

    # Interactive selection
    labels = [i.display_name for i in instances]
    print(prompt)
    for i, label in enumerate(labels):
        print(f"  [{i + 1}] {label}")
    if allow_all:
        print("  [a] All")

    try:
        choice = input("Choice: ").strip().lower()
    except (EOFError, KeyboardInterrupt):
        raise VastlyError("Cancelled.")

    if allow_all and choice == "a":
        return list(instances)

    if not choice.isdigit():
        raise VastlyError("Invalid choice.")

    idx = int(choice)
    if idx < 1 or idx > len(instances):
        raise VastlyError("Invalid choice.")

    return [instances[idx - 1]]
