#
#   Imandra Inc.
#
#   codelogician/modeling/dep_context.py
#

from __future__ import annotations

from collections import deque
from collections.abc import Iterable
from dataclasses import dataclass
from pathlib import Path
from typing import cast

from imandra.u.agents.code_logician.base import FormalizationDependency, ModuleInfo

from ..util import filter as cl_filter, map as cl_map
from .model import Model


@dataclass(frozen=True)
class CycleEdge:
    src: Model
    dst: Model


@dataclass(frozen=True)
class CycleInfo:
    """
    A concrete dependency cycle.

    - nodes is a closed walk: nodes[0] == nodes[-1]
    - edges align with nodes: edges[i] is nodes[i] -> nodes[i+1]
    """

    nodes: list[Model]
    edges: list[CycleEdge]

    @property
    def rel_paths(self) -> list[str]:
        return [m.rel_path for m in self.nodes]

    @property
    def edge_paths(self) -> list[tuple[str, str]]:
        return [(e.src.rel_path, e.dst.rel_path) for e in self.edges]


@dataclass(frozen=True)
class ModelDepContext:
    """
    Dependency analysis utilities for Model objects.

    This isolates graph traversal, dependency-to-ModuleInfo conversion,
    and "deps changed" computations away from Model and task policy.
    """

    @staticmethod
    def deps(model: Model) -> list[Model]:
        # "live" deps are Model objects; deserialized ones may be strings
        if not model.dependencies or isinstance(model.dependencies[0], str):
            return []
        return cast(list[Model], model.dependencies)

    @staticmethod
    def deps_paths(model: Model) -> list[str]:
        return [d.rel_path for d in ModelDepContext.deps(model)]

    @staticmethod
    def no_module_deps(model: Model) -> bool:
        return not model.dependencies or isinstance(model.dependencies[0], str)

    # -------------------------
    # ModuleInfo conversion
    # -------------------------
    @staticmethod
    def make_dep_name_from_path(rel_path: str) -> str:
        p = Path(rel_path).parent / Path(rel_path).stem
        return ('_'.join(Path(p).parts)).capitalize()

    @staticmethod
    def to_src_module_info(model: Model) -> ModuleInfo:
        return ModuleInfo(
            name=ModelDepContext.make_dep_name_from_path(model.rel_path),
            relative_path=Path(model.rel_path),
            content=model.src_code or '',
            src_lang=model.src_language,
        )

    @staticmethod
    def to_iml_module_info(model: Model) -> ModuleInfo | None:
        iml = model.iml_code()
        if iml is None:
            return None
        return ModuleInfo(
            name=ModelDepContext.make_dep_name_from_path(model.rel_path),
            relative_path=Path(model.rel_path),
            content=str(iml),
            src_lang='IML',
        )

    # -------------------------
    # Graph traversal / sorting
    # -------------------------
    @staticmethod
    def full_deep_dependencies(model: Model) -> list[Model]:
        result: list[Model] = []

        def visit(m: Model) -> None:
            for d in ModelDepContext.deps(m):
                if d not in result:
                    result.append(d)
                visit(d)

        visit(model)
        return result

    @staticmethod
    def topological_sort(models: list[Model]) -> list[Model]:
        """
        Return list sorted topologically by dependencies.
        (Same algorithm you had, but moved here.)
        """
        path_to_idx: dict[str, int] = {}
        idx_to_path: dict[int, str] = {}
        edges: list[tuple[int, int]] = []
        idx = 0

        def add_to_lookup(i: int, path: str) -> None:
            path_to_idx[path] = i
            idx_to_path[i] = path

        for m in models:
            if m.rel_path not in path_to_idx:
                add_to_lookup(idx, m.rel_path)
                idx += 1

            for dep_path in ModelDepContext.deps_paths(m):
                if dep_path not in path_to_idx:
                    add_to_lookup(idx, dep_path)
                    idx += 1
                edges.append((path_to_idx[m.rel_path], path_to_idx[dep_path]))

        def construct_adj(v: int, es: list[tuple[int, int]]) -> list[list[int]]:
            adj = [[] for _ in range(v)]
            for u, w in es:
                adj[u].append(w)
            return adj

        def topo(v: int, es: list[tuple[int, int]]) -> list[int]:
            adj = construct_adj(v, es)
            indegree = [0] * v
            for u in range(v):
                for w in adj[u]:
                    indegree[w] += 1
            q = deque([i for i in range(v) if indegree[i] == 0])

            out: list[int] = []
            while q:
                node = q.popleft()
                out.append(node)
                for nb in adj[node]:
                    indegree[nb] -= 1
                    if indegree[nb] == 0:
                        q.append(nb)

            if len(out) != v:
                raise Exception('Graph contains cycle!')
            return out

        model_by_path = {m.rel_path: m for m in models}
        sorted_idxs = topo(len(models), edges)
        return [model_by_path[idx_to_path[i]] for i in reversed(sorted_idxs)]

    # -------------------------
    # Cycle detection (returns actual nodes + actual edges)
    # -------------------------
    @staticmethod
    def detect_dependency_cycles(model: Model) -> list[CycleInfo]:
        """
        Detect directed cycles reachable from `model` following `dependencies`.

        Returns concrete CycleInfo (nodes + edges).

        Notes:
        - Designed for "live" graphs (dependencies as Model objects).
        - If deps are not live (strings / empty), returns [].
        """
        if ModelDepContext.no_module_deps(model):
            return []

        WHITE, GRAY, BLACK = 0, 1, 2  # noqa
        color: dict[str, int] = {}
        parent: dict[str, str | None] = {}
        node_by_path: dict[str, Model] = {}
        cycles: list[CycleInfo] = []

        def reconstruct_cycle_paths(from_path: str, to_path: str) -> list[str]:
            # Cycle is: to_path -> ... -> from_path -> to_path
            path = [to_path]
            cur = from_path
            while cur != to_path and cur is not None:
                path.append(cur)
                cur = parent.get(cur)
            path.append(to_path)
            path.reverse()
            return path

        def record_cycle(from_path: str, to_path: str) -> None:
            path = reconstruct_cycle_paths(from_path, to_path)

            # Convert to concrete nodes
            nodes: list[Model] = []
            for p in path:
                m = node_by_path.get(p)
                if m is None:
                    return
                nodes.append(m)

            # Build concrete edges and sanity-check adjacency
            edges: list[CycleEdge] = []
            for i in range(len(nodes) - 1):
                src = nodes[i]
                dst = nodes[i + 1]
                if dst not in ModelDepContext.deps(src):
                    # Graph mutated during traversal or inconsistent parent pointers.
                    return
                edges.append(CycleEdge(src=src, dst=dst))

            # Closed walk invariant (best-effort)
            if nodes and nodes[0] is nodes[-1]:
                cycles.append(CycleInfo(nodes=nodes, edges=edges))
            elif nodes and nodes[0].rel_path == nodes[-1].rel_path:
                # Fallback if identity differs but rel_path is same (shouldn’t happen in one graph)
                cycles.append(CycleInfo(nodes=nodes, edges=edges))

        def dfs(m: Model) -> None:
            mp = m.rel_path
            node_by_path[mp] = m
            color[mp] = GRAY

            for d in ModelDepContext.deps(m):
                dp = d.rel_path
                node_by_path[dp] = d

                if dp not in color:
                    parent[dp] = mp
                    dfs(d)
                elif color[dp] == GRAY:
                    # back-edge => cycle
                    record_cycle(mp, dp)

            color[mp] = BLACK

        parent[model.rel_path] = None
        dfs(model)
        return cycles

    @staticmethod
    def detect_cycles_in_models(models: Iterable[Model]) -> list[CycleInfo]:
        """
        Detect cycles anywhere in the given model set (live graph).
        """
        out: list[CycleInfo] = []
        seen_roots: set[str] = set()

        for m in models:
            if m.rel_path in seen_roots:
                continue
            seen_roots.add(m.rel_path)
            out.extend(ModelDepContext.detect_dependency_cycles(m))

        return out

    # -------------------------
    # Dependency payloads for CL
    # -------------------------
    @staticmethod
    def shallow_dep_iml_models(model: Model) -> list[FormalizationDependency]:
        if ModelDepContext.no_module_deps(model):
            return []

        ds: list[FormalizationDependency] = []
        for dep in ModelDepContext.deps(model):
            src_mod = ModelDepContext.to_src_module_info(dep)
            iml_mod = ModelDepContext.to_iml_module_info(dep)
            if iml_mod is None:
                continue
            ds.append(FormalizationDependency(src_module=src_mod, iml_module=iml_mod))
        return ds

    @staticmethod
    def dependencies_iml_models(model: Model) -> list[FormalizationDependency]:
        """
        Deep dependency list, topologically ordered, including only deps that have IML.
        """
        if ModelDepContext.no_module_deps(model):
            return []

        full_list = ModelDepContext.full_deep_dependencies(model)
        sorted_list = ModelDepContext.topological_sort(full_list)

        ds: list[FormalizationDependency] = []
        for dep in sorted_list:
            src_mod = ModelDepContext.to_src_module_info(dep)
            iml_mod = ModelDepContext.to_iml_module_info(dep)
            if iml_mod is None:
                continue
            ds.append(FormalizationDependency(src_module=src_mod, iml_module=iml_mod))
        return ds

    @staticmethod
    def _dep_fingerprint(d: FormalizationDependency) -> tuple[str, str]:
        """
        Stable fingerprint for dependency-change detection.

        We intentionally DO NOT include src_module.content because:
        - deps_changed is used to decide whether the last formalization is still valid
        - the invariant we care about is that imported IML for deps hasn't changed
        - src_module.content can change independently and would cause perpetual re-formalization

        Fingerprint = (relative_path, iml_content)
        """
        rel = str(d.src_module.relative_path)
        iml = d.iml_module.content if d.iml_module is not None else ''
        return (rel, iml)

    @staticmethod
    def deps_changed(model: Model) -> bool:
        """
        Compare current shallow deps (with IML) against deps recorded at last formalization.
        """
        current = ModelDepContext.shallow_dep_iml_models(model)

        paths_list = cl_map(Path, ModelDepContext.deps_paths(model))
        in_paths = lambda x: x.src_module.relative_path in paths_list  # noqa: E731
        relevant_formalized = cl_filter(in_paths, model.formalized_deps)

        current_fp = {ModelDepContext._dep_fingerprint(d) for d in current}
        formalized_fp = {
            ModelDepContext._dep_fingerprint(d) for d in relevant_formalized
        }

        return current_fp != formalized_fp

    @staticmethod
    def create_import_iml_statements(model: Model) -> str:
        text = ''
        for d in ModelDepContext.dependencies_iml_models(model):
            name = d.iml_module.name
            rel_path = str(Path(d.iml_module.relative_path).with_suffix('.iml'))
            text += f'[@@@import {name}, "{rel_path}"]\n'
        return text

    @staticmethod
    def to_CL_src_module_info(m: Model) -> ModuleInfo:
        """
        Create ModuleInfo
        """
        return ModuleInfo(
            name=ModelDepContext.make_dep_name_from_path(m.rel_path),
            relative_path=Path(m.rel_path),
            content=m.src_code if m.src_code else '',
            src_lang=m.src_language,
        )

    @staticmethod
    def to_CL_iml_module_info(m: Model) -> ModuleInfo | None:
        """
        Return ModuleInfo object for the IML code if it's available, None otherwise
        """

        # If we don't have any IML code for this model, let's not return it
        if m.iml_code() is None:
            return None

        if m.iml_code() is not None:
            iml_code = str(m.iml_code())
        else:
            iml_code = ''

        return ModuleInfo(
            name=ModelDepContext.make_dep_name_from_path(m.rel_path),
            relative_path=Path(m.rel_path),
            content=iml_code,
            src_lang='IML',
        )
