#
#   Imandra Inc.
#
#   speclogician/demos/loader.py
#

from __future__ import annotations

import importlib.util
from pathlib import Path
from typing import Any, Callable

from speclogician.demos.model import DemoSpec
from speclogician.state.state import State
from speclogician.state.change import StateChange


# -----------------------------------------------------------------------------
# discovery / selection
# -----------------------------------------------------------------------------

def discover_demos(demos_dir: Path) -> list[DemoSpec]:
    demos: list[DemoSpec] = []
    if not demos_dir.exists():
        return demos

    for p in sorted(demos_dir.iterdir()):
        if not p.is_dir():
            continue
        if (p / "changes.py").exists():
            demos.append(DemoSpec(name=p.name, root=p))
    return demos


def load_demo(demos_dir: Path, name: str) -> DemoSpec:
    root = demos_dir / name
    if not root.exists() or not root.is_dir():
        raise FileNotFoundError(f'demo not found: {name}')
    if not (root / "changes.py").exists():
        raise FileNotFoundError('demo missing changes.py')
    return DemoSpec(name=name, root=root)


# -----------------------------------------------------------------------------
# changes.py loader
# -----------------------------------------------------------------------------

def _load_module_from_path(py_path: Path, module_name: str) -> Any:
    spec = importlib.util.spec_from_file_location(module_name, str(py_path))
    if spec is None or spec.loader is None:
        raise RuntimeError(f"Could not import: {py_path}")
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)  # type: ignore[attr-defined]
    return mod


def load_changes(demo: DemoSpec) -> list[StateChange]:
    """
    Returns the list in `CHANGES` from the demo's changes.py.
    Kept as `Any` to avoid importing all StateChange subclasses here.
    """
    mod = _load_module_from_path(
        demo.changes_path,
        module_name=f"speclogician_demo_{demo.name}",
    )
    if not hasattr(mod, "CHANGES"):
        raise RuntimeError(f"{demo.changes_path} must define CHANGES = [...]")
    changes = getattr(mod, "CHANGES")
    if not isinstance(changes, list):
        raise RuntimeError(f"{demo.changes_path}: CHANGES must be a list")
    return changes


def load_demo_name(demo: DemoSpec) -> str:
    mod = _load_module_from_path(
        demo.changes_path,
        module_name=f"speclogician_demo_name_{demo.name}",
    )
    return getattr(mod, "DEMO_NAME", demo.name)


# -----------------------------------------------------------------------------
# cache (state.json in demo folder)
# -----------------------------------------------------------------------------

def load_cached_state(demo: DemoSpec) -> State:
    """
    Used by `speclogician demos tui <name>` and `... shell <name>`.

    Strict: requires demo.state.json to exist.
    """
    p = demo.cache_path
    if not p.exists():
        raise FileNotFoundError(
            f"Demo cache missing: {p}\n"
            f"Expected a JSON-serialized State model at {p.name}."
        )
    return State.model_validate_json(p.read_text(encoding="utf-8"))


def save_cached_state(demo: DemoSpec, state: State) -> None:
    """
    Writes demo.state.json (State model JSON) into the demo directory.
    """
    demo.root.mkdir(parents=True, exist_ok=True)
    demo.cache_path.write_text(
        state.model_dump_json(indent=2),
        encoding="utf-8",
    )


# -----------------------------------------------------------------------------
# execution (always run changes from empty state)
# -----------------------------------------------------------------------------

ApplyChangeFn = Callable[[State, Any], State]


def run_changes_fresh(demo: DemoSpec, *, apply_change: ApplyChangeFn) -> State:
    """
    Used by `speclogician demos run <name>`.

    Always starts from an empty State and applies every change from changes.py.

    `apply_change` is injected to keep this module independent of the engine
    implementation. Typical usage:

        def apply_change(state: State, ch: StateChange) -> State:
            return state.process_change(ch)

    """
    state = State()
    for ch in load_changes(demo):
        state = apply_change(state, ch)
    return state