"""Command + processor registry (YAML-driven).

This module is the foundation for decoupling validation:
- execution command definitions (how to run a step)
- output processors (how to summarize/structure outputs)

The registry can be loaded from a local YAML file and optionally extended via "includes"
that point to other YAML files. This enables storing command/processor definitions outside
of the aida-mcp package/repo while keeping a stable configuration contract.
"""

from __future__ import annotations

import logging
import os
from collections.abc import Mapping, MutableMapping
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, cast

import yaml

from aida_mcp.utils.aida_paths import DOT_AIDA

logger = logging.getLogger(__name__)

ENV_VALIDATION_REGISTRY_PATH = "AIDA_VALIDATION_REGISTRY_PATH"
ENV_WORKSPACE_ROOT = "AIDA_WORKSPACE_ROOT"


def _workspace_root_from_env() -> Path | None:
    # Keep this sync and conservative: registry load must not depend on MCP Context.
    override = os.environ.get(ENV_WORKSPACE_ROOT)
    if not override:
        return None
    try:
        candidate = Path(override).expanduser()
        if not candidate.is_absolute():
            candidate = (Path.cwd() / candidate).resolve()
        return candidate if candidate.exists() else None
    except Exception:
        return None


def _default_registry_path() -> Path:
    """
    Resolve the default validation registry path.

    Priority:
    1) ENV override (`AIDA_VALIDATION_REGISTRY_PATH`)
    2) Workspace-root-based locations
    3) Working-directory-based locations
    4) Code-relative fallback (`config/validation_registry.yaml`)
    """
    env_path = os.environ.get(ENV_VALIDATION_REGISTRY_PATH)
    if isinstance(env_path, str) and env_path.strip():
        p = Path(env_path.strip()).expanduser()
        if not p.is_absolute():
            p = (Path.cwd() / p).resolve()
        return p

    ws = _workspace_root_from_env()
    if ws is not None:
        # Preferred repo-owned location.
        cand0 = ws / DOT_AIDA / "validation_registry.yaml"
        if cand0.exists():
            return cand0

        # Allow colocated config at workspace root.
        cand2 = ws / "config" / "validation_registry.yaml"
        if cand2.exists():
            return cand2

    cwd = Path.cwd().resolve()
    cand3 = cwd / DOT_AIDA / "validation_registry.yaml"
    if cand3.exists():
        return cand3
    cand4 = cwd / "config" / "validation_registry.yaml"
    if cand4.exists():
        return cand4

    # Final fallback: resolve relative to current package location.
    for parent in Path(__file__).resolve().parents:
        cand5 = parent / "config" / "validation_registry.yaml"
        if cand5.exists():
            return cand5
    return cwd / "config" / "validation_registry.yaml"


def _new_str_dict() -> dict[str, str]:
    # Keep this as a named factory (instead of `dict` / `lambda: {}`) to help Pyright
    # preserve the concrete `dict[str, str]` type for dataclass fields under strict settings.
    return {}


def _string_mapping(raw: object) -> dict[str, Any]:
    if not isinstance(raw, Mapping):
        return {}
    out: dict[str, Any] = {}
    m = cast(Mapping[object, object], raw)
    for k, v in m.items():
        out[str(k)] = v
    return out


def _string_list(raw: object) -> list[str]:
    if not isinstance(raw, list):
        return []
    items = cast(list[object], raw)
    if not all(isinstance(x, str) for x in items):
        return []
    return cast(list[str], items)


@dataclass(frozen=True, slots=True)
class CommandSpec:
    """Executable external command definition.

    Commands are always executed as external subprocesses with raw argv (no shell).
    Placeholder expansion is handled by callers.
    """

    argv: list[str]
    cwd: str | None = None
    env: dict[str, str] = field(default_factory=_new_str_dict)
    timeout_sec: int | None = None
    # Streaming is preferred (default) so we can forward progress updates to MCP clients.
    # Set to false per-command to force non-streaming execution.
    streaming: bool = True


@dataclass(frozen=True, slots=True)
class ProcessorSpec:
    """Output processor definition.

    Supported kinds (initial):
    - builtin: built-in processor shipped with aida-mcp
    - python_module: run `<sys.executable> -m <module> ...`
    - exec: run an external executable (argv[0]) directly
    """

    kind: str
    builtin_id: str | None = None
    module: str | None = None
    argv: list[str] | None = None
    timeout_sec: int | None = None
    env: dict[str, str] = field(default_factory=_new_str_dict)


@dataclass(frozen=True, slots=True)
class ValidationRegistry:
    """Registry of command and processor definitions."""

    commands: Mapping[str, CommandSpec]
    processors: Mapping[str, ProcessorSpec]

    @classmethod
    def load(cls, config_path: Path | None = None) -> ValidationRegistry:
        if config_path is None:
            config_path = _default_registry_path()

        if not config_path.exists():
            logger.info("Validation registry config not found: %s", config_path)
            return cls(commands={}, processors={})

        try:
            reg = _load_registry_with_includes(config_path)
            return reg
        except Exception as exc:
            logger.warning("Failed to load validation registry from %s: %s", config_path, exc)
            return cls(commands={}, processors={})


def _load_registry_with_includes(config_path: Path) -> ValidationRegistry:
    """Load a registry YAML file with optional includes.

    Merge semantics:
    - Included files are loaded first (in order), then the main file is overlaid last.
    - Later definitions override earlier ones by id.
    """

    seen: set[Path] = set()
    commands: MutableMapping[str, CommandSpec] = {}
    processors: MutableMapping[str, ProcessorSpec] = {}

    def _load_one(path: Path) -> None:
        resolved = path.resolve()
        if resolved in seen:
            return
        seen.add(resolved)

        try:
            raw_text = resolved.read_text(encoding="utf-8")
            loaded: object = yaml.safe_load(raw_text)
        except Exception as exc:
            logger.warning("Skipping registry file %s (read/parse error): %s", resolved, exc)
            return

        root = _string_mapping(loaded)
        reg_any = root.get("registry")
        reg = _string_mapping(reg_any)

        base_dir = resolved.parent
        includes = _string_list(reg.get("includes"))
        for inc in includes:
            if not inc:
                continue
            inc_path = Path(inc)
            if not inc_path.is_absolute():
                inc_path = (base_dir / inc_path).resolve()
            _load_one(inc_path)

        commands_any = reg.get("commands")
        if isinstance(commands_any, Mapping):
            for cid, cval in _string_mapping(cast(object, commands_any)).items():
                spec = _parse_command_spec(cval)
                if spec is not None:
                    commands[cid] = spec

        processors_any = reg.get("processors")
        if isinstance(processors_any, Mapping):
            for pid, pval in _string_mapping(cast(object, processors_any)).items():
                spec = _parse_processor_spec(pval)
                if spec is not None:
                    processors[pid] = spec

    # Includes get loaded inside _load_one; we still want overlay-last behavior.
    _load_one(config_path)
    return ValidationRegistry(commands=commands, processors=processors)


def _parse_command_spec(raw: object) -> CommandSpec | None:
    if not isinstance(raw, Mapping):
        return None
    m = _string_mapping(cast(object, raw))

    argv = _string_list(m.get("argv"))
    if not argv:
        return None

    cwd_any = m.get("cwd")
    cwd = cwd_any if isinstance(cwd_any, str) and cwd_any else None

    env: dict[str, str] = {}
    env_any = m.get("env")
    if isinstance(env_any, Mapping):
        env_map = cast(Mapping[object, object], env_any)
        for k, v in env_map.items():
            if isinstance(v, str):
                env[str(k)] = v

    timeout_any = m.get("timeout_sec")
    timeout_sec = timeout_any if isinstance(timeout_any, int) and timeout_any > 0 else None

    streaming_any = m.get("streaming")
    # Breaking-change default: streaming is ON unless explicitly disabled.
    streaming = streaming_any if isinstance(streaming_any, bool) else True

    return CommandSpec(
        argv=argv,
        cwd=cwd,
        env=env,
        timeout_sec=timeout_sec,
        streaming=streaming,
    )


def _parse_processor_spec(raw: object) -> ProcessorSpec | None:
    if not isinstance(raw, Mapping):
        return None
    m = _string_mapping(cast(object, raw))

    module_any = m.get("module")
    module = module_any if isinstance(module_any, str) and module_any else None

    builtin_id_any = m.get("builtin_id")
    builtin_id = builtin_id_any if isinstance(builtin_id_any, str) and builtin_id_any.strip() else None

    argv_list = _string_list(m.get("argv"))
    argv = argv_list if argv_list else None

    timeout_any = m.get("timeout_sec")
    timeout_sec = timeout_any if isinstance(timeout_any, int) and timeout_any > 0 else None

    env: dict[str, str] = {}
    env_any = m.get("env")
    if isinstance(env_any, Mapping):
        env_map = cast(Mapping[object, object], env_any)
        for k, v in env_map.items():
            if isinstance(v, str):
                env[str(k)] = v

    # Unified validate supports builtin processors only.
    kind_any = m.get("kind")
    kind = kind_any if isinstance(kind_any, str) and kind_any else "builtin"
    if kind != "builtin" or not builtin_id:
        return None

    return ProcessorSpec(kind=kind, builtin_id=builtin_id, module=module, argv=argv, timeout_sec=timeout_sec, env=env)


_cached_registry: ValidationRegistry | None = None
_cached_registry_path: Path | None = None
_cached_registry_mtime: float | None = None


def get_validation_registry() -> ValidationRegistry:
    """Get (cached) validation command/processor registry."""
    global _cached_registry, _cached_registry_mtime, _cached_registry_path

    config_path = _default_registry_path()
    try:
        mtime = config_path.stat().st_mtime if config_path.exists() else None
    except Exception:
        mtime = None

    if _cached_registry is None or _cached_registry_path != config_path or _cached_registry_mtime != mtime:
        _cached_registry = ValidationRegistry.load(config_path=config_path)
        _cached_registry_path = config_path
        _cached_registry_mtime = mtime

    return _cached_registry
