from __future__ import annotations

import argparse
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence, Set, Tuple, Union

import yaml


# Make sure the repository root is importable so we can import
# `sync_state_machine.common.types` when running as a script.
_REPO_ROOT = Path(__file__).resolve().parents[1]
if str(_REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(_REPO_ROOT))


class ValidationError(Exception):
    pass


@dataclass(frozen=True)
class Problem:
    level: str  # ERROR | WARN
    message: str
    path: str


def _load_yaml(path: Path) -> Dict[str, Any]:
    try:
        with path.open("r", encoding="utf-8") as f:
            data = yaml.safe_load(f)
    except Exception as exc:
        raise ValidationError(f"Failed to read YAML: {path}: {exc}") from exc

    if not isinstance(data, dict):
        raise ValidationError(f"Top-level YAML must be a mapping/dict: {path}")
    return data


def _as_list(x: Any) -> List[Any]:
    if x is None:
        return []
    if isinstance(x, list):
        return x
    return [x]


def _require_keys(obj: Mapping[str, Any], keys: Sequence[str], base_path: str, problems: List[Problem]) -> None:
    for k in keys:
        if k not in obj:
            problems.append(Problem("ERROR", f"Missing required key '{k}'", f"{base_path}.{k}"))


def _import_enums() -> Tuple[Set[str], Set[str], Set[str]]:
    # Only validate enum NAMES; values (lowercase strings) are implementation details.
    from sync_state_machine.common.types import BindingStatus, SyncAction, SyncStatus

    binding = {e.name for e in BindingStatus}
    action = {e.name for e in SyncAction}
    status = {e.name for e in SyncStatus}
    return binding, action, status


def _validate_state_def(
    state_id: str,
    state_def: Mapping[str, Any],
    binding_names: Set[str],
    action_names: Set[str],
    status_names: Set[str],
    problems: List[Problem],
) -> None:
    base = f"states.{state_id}"
    _require_keys(state_def, ["binding_status", "action", "status", "desc"], base, problems)

    def _validate_enum_or_list(
        value: Any,
        *,
        allowed: Set[str],
        field_path: str,
        label: str,
    ) -> None:
        if isinstance(value, str):
            if value != "*" and value not in allowed:
                problems.append(Problem("ERROR", f"Unknown {label} name: {value}", field_path))
            return
        if isinstance(value, list):
            if not value:
                problems.append(Problem("ERROR", f"{label} list must be non-empty", field_path))
                return
            for i, item in enumerate(value):
                if not isinstance(item, str):
                    problems.append(Problem("ERROR", f"{label} list items must be strings", f"{field_path}[{i}]"))
                    continue
                if item != "*" and item not in allowed:
                    problems.append(Problem("ERROR", f"Unknown {label} name: {item}", f"{field_path}[{i}]"))
            return
        problems.append(Problem("ERROR", f"{label} must be a string, '*' or a list of strings", field_path))

    _validate_enum_or_list(
        state_def.get("binding_status"),
        allowed=binding_names,
        field_path=f"{base}.binding_status",
        label="BindingStatus",
    )
    _validate_enum_or_list(
        state_def.get("action"),
        allowed=action_names,
        field_path=f"{base}.action",
        label="SyncAction",
    )
    _validate_enum_or_list(
        state_def.get("status"),
        allowed=status_names,
        field_path=f"{base}.status",
        label="SyncStatus",
    )

    # Optional: whether the node has data_id at this state.
    # Used for signature uniqueness checks; kept optional to avoid forcing all configs to specify it.
    data_id_exist = state_def.get("data_id_exist", None)
    if data_id_exist is not None:
        field_path = f"{base}.data_id_exist"
        if data_id_exist == "*":
            pass
        elif isinstance(data_id_exist, bool):
            pass
        elif isinstance(data_id_exist, list):
            if not data_id_exist:
                problems.append(Problem("ERROR", "data_id_exist list must be non-empty", field_path))
            else:
                for i, item in enumerate(data_id_exist):
                    if item == "*":
                        continue
                    if not isinstance(item, bool):
                        problems.append(
                            Problem(
                                "ERROR",
                                "data_id_exist list items must be bool or '*'",
                                f"{field_path}[{i}]",
                            )
                        )
        else:
            problems.append(Problem("ERROR", "data_id_exist must be bool, '*', or a list of them", field_path))


def _expand_enum_values(value: Any, *, all_values: Set[str]) -> Set[str]:
    if isinstance(value, str):
        return set(all_values) if value == "*" else {value}
    if isinstance(value, list):
        out: Set[str] = set()
        for item in value:
            if not isinstance(item, str):
                continue
            if item == "*":
                out |= set(all_values)
            else:
                out.add(item)
        return out
    return set()


def _expand_data_id_exist(value: Any) -> Set[bool]:
    # Missing is treated as wildcard for uniqueness check (strict).
    if value is None or value == "*":
        return {True, False}
    if isinstance(value, bool):
        return {value}
    if isinstance(value, list):
        out: Set[bool] = set()
        for item in value:
            if item == "*":
                out |= {True, False}
            elif isinstance(item, bool):
                out.add(item)
        return out if out else {True, False}
    return {True, False}


def _validate_state_signature_uniqueness(
    states: Mapping[str, Any],
    *,
    binding_names: Set[str],
    action_names: Set[str],
    status_names: Set[str],
    problems: List[Problem],
) -> None:
    """Ensure no two states can share the same (binding_status, action, status, data_id_exist).

    Semantics:
    - '*' expands to all enum values.
    - list expands to the set.
    - data_id_exist missing expands to {True, False} (strict), so uniqueness failures surface.
    """

    sig_to_states: Dict[Tuple[str, str, str, bool], List[str]] = {}

    for sid, sdef in states.items():
        if not isinstance(sdef, dict):
            continue

        binding_set = _expand_enum_values(sdef.get("binding_status"), all_values=binding_names)
        action_set = _expand_enum_values(sdef.get("action"), all_values=action_names)
        status_set = _expand_enum_values(sdef.get("status"), all_values=status_names)
        data_id_set = _expand_data_id_exist(sdef.get("data_id_exist"))

        # If any of the sets are empty due to prior validation errors, skip to avoid noise.
        if not binding_set or not action_set or not status_set:
            continue

        for b in binding_set:
            for a in action_set:
                for st in status_set:
                    for de in data_id_set:
                        sig_to_states.setdefault((b, a, st, de), []).append(str(sid))

    collisions = [(sig, sids) for sig, sids in sig_to_states.items() if len(set(sids)) > 1]
    collisions.sort(key=lambda x: (-len(set(x[1])), x[0]))

    for sig, sids in collisions:
        uniq = sorted(set(sids))
        problems.append(
            Problem(
                "ERROR",
                f"State signature overlaps for (binding_status, action, status, data_id_exist)={sig}: {', '.join(uniq)}",
                "states",
            )
        )


def _collect_state_ids(cfg: Mapping[str, Any], problems: List[Problem]) -> Set[str]:
    states = cfg.get("states")
    if not isinstance(states, dict):
        problems.append(Problem("ERROR", "'states' must be a mapping", "states"))
        return set()
    return set(states.keys())


def _validate_context(cfg: Mapping[str, Any], problems: List[Problem]) -> Set[str]:
    ctx = cfg.get("context")
    if not isinstance(ctx, dict):
        problems.append(Problem("ERROR", "'context' must be a mapping", "context"))
        return set()

    ctx_keys: Set[str] = set()
    for k, v in ctx.items():
        ctx_keys.add(k)
        if not isinstance(v, dict):
            problems.append(Problem("ERROR", "context entry must be a mapping", f"context.{k}"))
            continue
        if "type" not in v:
            problems.append(Problem("ERROR", "context entry missing 'type'", f"context.{k}.type"))
        if "desc" not in v:
            problems.append(Problem("WARN", "context entry missing 'desc'", f"context.{k}.desc"))
        if v.get("type") == "enum":
            vals = v.get("values")
            if not isinstance(vals, list) or not vals:
                problems.append(Problem("ERROR", "enum context must have non-empty 'values'", f"context.{k}.values"))
    return ctx_keys


def _validate_stage(cfg: Mapping[str, Any], stage_id: str, stage_def: Mapping[str, Any], state_ids: Set[str], problems: List[Problem]) -> None:
    base = f"stages.{stage_id}"
    _require_keys(stage_def, ["desc", "entry_states", "exit_states"], base, problems)

    for key in ("entry_states", "exit_states"):
        vals = stage_def.get(key)
        if not isinstance(vals, list) or not vals:
            problems.append(Problem("ERROR", f"{key} must be a non-empty list", f"{base}.{key}"))
            continue
        for s in vals:
            if s == "*":
                continue
            if s not in state_ids and s not in (cfg.get("pseudo_states") or {}):
                problems.append(Problem("ERROR", f"Unknown state id '{s}'", f"{base}.{key}"))


def _validate_stages(cfg: Mapping[str, Any], state_ids: Set[str], problems: List[Problem]) -> Set[str]:
    stages = cfg.get("stages")
    if not isinstance(stages, dict):
        problems.append(Problem("ERROR", "'stages' must be a mapping", "stages"))
        return set()

    stage_ids: Set[str] = set()
    for sid, sdef in stages.items():
        stage_ids.add(sid)
        if not isinstance(sdef, dict):
            problems.append(Problem("ERROR", "stage definition must be a mapping", f"stages.{sid}"))
            continue
        _validate_stage(cfg, sid, sdef, state_ids, problems)
    return stage_ids


def _guard_keys(guard: Mapping[str, Any]) -> Set[str]:
    keys: Set[str] = set()
    for k, v in guard.items():
        if isinstance(v, dict):
            # Nested structures are not allowed for now.
            continue
        keys.add(k)
    return keys


def _validate_transition(
    tid: str,
    tdef: Mapping[str, Any],
    stage_ids: Set[str],
    state_ids: Set[str],
    pseudo_state_ids: Set[str],
    context_keys: Set[str],
    problems: List[Problem],
) -> None:
    base = f"transitions.{tid}"
    _require_keys(tdef, ["stage", "desc", "from", "outcomes"], base, problems)

    stage = tdef.get("stage")
    if isinstance(stage, str) and stage not in stage_ids:
        problems.append(Problem("ERROR", f"Unknown stage '{stage}'", f"{base}.stage"))

    frm = tdef.get("from")
    if not isinstance(frm, list) or not frm:
        problems.append(Problem("ERROR", "'from' must be a non-empty list", f"{base}.from"))
    else:
        for s in frm:
            if s == "*":
                continue
            if s not in state_ids and s not in pseudo_state_ids:
                problems.append(Problem("ERROR", f"Unknown from-state '{s}'", f"{base}.from"))

    outcomes = tdef.get("outcomes")
    if not isinstance(outcomes, list) or not outcomes:
        problems.append(Problem("ERROR", "'outcomes' must be a non-empty list", f"{base}.outcomes"))
        return

    # Basic overlap check for outcome guards.
    normalized_guards: List[Tuple[int, Dict[str, Any]]] = []

    for i, o in enumerate(outcomes):
        opath = f"{base}.outcomes[{i}]"
        if not isinstance(o, dict):
            problems.append(Problem("ERROR", "outcome must be a mapping", opath))
            continue
        if "when" not in o:
            problems.append(Problem("ERROR", "outcome missing 'when'", f"{opath}.when"))
            continue
        when = o.get("when")
        if not isinstance(when, dict):
            problems.append(Problem("ERROR", "'when' must be a mapping", f"{opath}.when"))
            continue

        # Guard variable key validation
        for k in _guard_keys(when):
            if k not in context_keys and k not in ("self.status",):
                problems.append(Problem("ERROR", f"Unknown context key '{k}'", f"{opath}.when.{k}"))

        has_to = "to" in o
        has_to_map = "to_map" in o
        if not (has_to or has_to_map):
            problems.append(Problem("ERROR", "outcome must have 'to' or 'to_map'", opath))
            continue

        if has_to:
            to = o.get("to")
            if not isinstance(to, str):
                problems.append(Problem("ERROR", "'to' must be a state id string", f"{opath}.to"))
            else:
                if to not in state_ids and to not in pseudo_state_ids:
                    problems.append(Problem("ERROR", f"Unknown to-state '{to}'", f"{opath}.to"))

        if has_to_map:
            tm = o.get("to_map")
            if not isinstance(tm, dict):
                problems.append(Problem("ERROR", "'to_map' must be a mapping", f"{opath}.to_map"))
            else:
                if tm.get("key") != "self.status":
                    problems.append(Problem("WARN", "Only to_map.key=self.status is supported by validator", f"{opath}.to_map.key"))
                mp = tm.get("map")
                if not isinstance(mp, dict) or not mp:
                    problems.append(Problem("ERROR", "to_map.map must be a non-empty mapping", f"{opath}.to_map.map"))
                else:
                    for k, v in mp.items():
                        if not isinstance(k, str):
                            problems.append(Problem("ERROR", "to_map.map keys must be strings", f"{opath}.to_map.map"))
                        if not isinstance(v, str) or v not in state_ids:
                            problems.append(Problem("ERROR", f"Unknown to_map target '{v}'", f"{opath}.to_map.map.{k}"))

        emit = o.get("emit")
        if emit is not None:
            if not isinstance(emit, dict):
                problems.append(Problem("ERROR", "emit must be a mapping", f"{opath}.emit"))
            else:
                if "spawn_target_node_state" in emit:
                    spawn = emit.get("spawn_target_node_state")
                    if spawn not in state_ids:
                        problems.append(Problem("ERROR", f"Unknown spawn_target_node_state '{spawn}'", f"{opath}.emit.spawn_target_node_state"))

        normalized_guards.append((i, dict(when)))

    # Overlap check: if two guards are compatible (no contradictory assignments), they overlap.
    for i in range(len(normalized_guards)):
        idx_a, ga = normalized_guards[i]
        for j in range(i + 1, len(normalized_guards)):
            idx_b, gb = normalized_guards[j]
            if _guards_overlap(ga, gb):
                problems.append(
                    Problem(
                        "WARN",
                        f"Outcome guards may overlap (order-dependent): outcomes[{idx_a}] vs outcomes[{idx_b}]",
                        f"{base}.outcomes",
                    )
                )


def _guards_overlap(a: Mapping[str, Any], b: Mapping[str, Any]) -> bool:
    for k, va in a.items():
        if k not in b:
            continue
        vb = b[k]
        if va != vb:
            return False
    for k, vb in b.items():
        if k not in a:
            continue
        va = a[k]
        if va != vb:
            return False
    return True


def validate_config(cfg_path: Path) -> Tuple[List[Problem], Dict[str, Any]]:
    cfg = _load_yaml(cfg_path)

    problems: List[Problem] = []

    binding_names, action_names, status_names = _import_enums()

    # meta
    meta = cfg.get("meta")
    if not isinstance(meta, dict):
        problems.append(Problem("ERROR", "'meta' must be a mapping", "meta"))

    # context
    context_keys = _validate_context(cfg, problems)

    # states
    state_ids = _collect_state_ids(cfg, problems)
    states = cfg.get("states") if isinstance(cfg.get("states"), dict) else {}
    for sid, sdef in states.items():
        if not isinstance(sdef, dict):
            problems.append(Problem("ERROR", "state definition must be a mapping", f"states.{sid}"))
            continue
        _validate_state_def(sid, sdef, binding_names, action_names, status_names, problems)

    # Ensure states are uniquely identifiable by a strict field signature.
    _validate_state_signature_uniqueness(
        states,
        binding_names=binding_names,
        action_names=action_names,
        status_names=status_names,
        problems=problems,
    )

    pseudo_states = cfg.get("pseudo_states")
    pseudo_state_ids: Set[str] = set()
    if pseudo_states is not None:
        if not isinstance(pseudo_states, dict):
            problems.append(Problem("ERROR", "'pseudo_states' must be a mapping", "pseudo_states"))
        else:
            pseudo_state_ids = set(pseudo_states.keys())

    # stages
    stage_ids = _validate_stages(cfg, state_ids, problems)

    # transitions
    transitions = cfg.get("transitions")
    if not isinstance(transitions, dict):
        problems.append(Problem("ERROR", "'transitions' must be a mapping", "transitions"))
    else:
        for tid, tdef in transitions.items():
            if not isinstance(tdef, dict):
                problems.append(Problem("ERROR", "transition definition must be a mapping", f"transitions.{tid}"))
                continue
            _validate_transition(tid, tdef, stage_ids, state_ids, pseudo_state_ids, context_keys, problems)

    # invariants: basic structural checks only (semantics validated by runtime engine)
    inv = cfg.get("invariants")
    if inv is not None and not isinstance(inv, dict):
        problems.append(Problem("ERROR", "'invariants' must be a mapping", "invariants"))

    return problems, cfg


def main(argv: Optional[Sequence[str]] = None) -> int:
    parser = argparse.ArgumentParser(description="Validate state machine YAML config")
    parser.add_argument(
        "--config",
        default=str(Path(__file__).resolve().parents[1] / "sync_state_machine" / "config" / "node_state_machine.yaml"),
        help="Path to YAML config",
    )
    args = parser.parse_args(list(argv) if argv is not None else None)

    cfg_path = Path(args.config)
    problems, _ = validate_config(cfg_path)

    errors = [p for p in problems if p.level == "ERROR"]
    warns = [p for p in problems if p.level == "WARN"]

    for p in problems:
        print(f"[{p.level}] {p.path}: {p.message}")

    print("\nSummary:")
    print(f"  errors: {len(errors)}")
    print(f"  warns : {len(warns)}")

    if errors:
        return 2
    if warns:
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())