#
#   Imandra Inc.
#
#   speclogician/state/recommender.py
#

from __future__ import annotations

from typing import Iterable, TypeAlias, Callable, Optional

from pydantic import BaseModel

from .diff import ComparisonOutcome, ValueDiff, StateDiff
from .inst import StateInstance
from .recommendation import Recommendation

class ChangeFailure(BaseModel):
    """
    Lightweight failure context for recommendation generation.
    This replaces dependency on ProcessChangeResult.
    """

    change: str                 # e.g. "AddPredicate", "EditScenario"
    stage: str                  # e.g. "apply_change", "analyze_change"
    error_type: str             # e.g. "TypeError"
    error: str                  # str(exception)

    # Optional extras if you want later
    message: Optional[str] = None


Rule: TypeAlias = Callable[[StateInstance, StateInstance, Optional[ChangeFailure]], Iterable[Recommendation]]


class Recommender:
    """
    Turn (old_state, new_state) into human-readable next steps.

    Typical use:
      recs = Recommender().recommend(old_inst, new_inst)
      # feed "\n".join(r.text for r in recs) into your LLM CLI loop
    """

    def __init__(self, rules: list[Rule] | None = None) -> None:
        self._rules: list[Rule] = rules or [
            _rule_error_triage,
            _rule_state_diff_summary,
            _rule_missing_components,
            _rule_inconsistent_scenarios,
            _rule_conflicting_scenarios,
            _rule_unmatched_traces,
            _rule_invalid_iml_components,
            _rule_traceability_nudge,
        ]

    def recommend(
        self,
        old_inst: StateInstance,
        new_inst: StateInstance,
        *,
        failure: ChangeFailure | None = None,
    ) -> list[Recommendation]:
        out: list[Recommendation] = []

        for rule in self._rules:
            try:
                out.extend(list(rule(old_inst, new_inst, failure)))
            except Exception as e:
                # recommender should never take down the CLI
                out.append(
                    Recommendation(
                        text=f"Recommender rule failed ({getattr(rule, '__name__', type(rule).__name__)}): {e}",
                        kind="warning",
                        priority=5,
                    )
                )

        # Deduplicate by `text` while preserving order.
        # If the same text appears multiple times, keep the one with the lowest priority number.
        chosen: dict[str, Recommendation] = {}
        order: list[str] = []

        for r in out:
            if r.text not in chosen:
                chosen[r.text] = r
                order.append(r.text)
            else:
                prev = chosen[r.text]
                if r.priority < prev.priority:
                    chosen[r.text] = r

        return [chosen[t] for t in order]


# ----------------------------
# Helpers
# ----------------------------

def _ok(failure: ChangeFailure | None) -> bool:
    """
    Mirror old `res["ok"]` meaning:
    - failure is None => ok
    - failure provided => failed
    """
    return failure is None


def _outcome(v: ValueDiff[object]) -> ComparisonOutcome:
    """
    Use the comp_func if present, otherwise fallback to raw equality.
    """
    if v.comp_func is not None:  # type: ignore[truthy-function]
        return v.comp_func(v.before, v.after)  # type: ignore[misc]
    return ComparisonOutcome.NO_CHANGE if v.before == v.after else ComparisonOutcome.UNKNOWN


def _rec(
    text: str,
    *,
    kind: str = "next",
    priority: int = 50,
) -> Recommendation:
    return Recommendation(text=text, kind=kind, priority=priority)  # type: ignore[arg-type]


# ----------------------------
# Rules (keep them small + boring)
# ----------------------------

def _rule_error_triage(
    old_inst: StateInstance,
    new_inst: StateInstance,
    failure: ChangeFailure | None,
) -> Iterable[Recommendation]:
    if _ok(failure):
        return []

    # If you have the actual StateChange object at the call site,
    # you can set failure.change = ch.__class__.__name__.
    return [
        Recommendation(
            kind="warning",
            priority=0,
            text=(
                f"The last change failed during **{failure.stage}** for `{failure.change}`.\n"
                f"Error: `{failure.error_type}: {failure.error}`\n"
                "Recommendation: fix this error first (no later changes will be reliable)."
            ),
        )
    ]


def _rule_missing_components(
    old_inst: StateInstance,
    new_inst: StateInstance,
    failure: ChangeFailure | None,
) -> Iterable[Recommendation]:
    if not _ok(failure):
        return []

    spec = new_inst.spec
    missing = [
        sc
        for sc in spec.scenarios
        if getattr(sc, "preds_missing", []) or getattr(sc, "trans_missing", [])
    ]
    if not missing:
        return []

    examples = missing[:3]
    bits: list[str] = []
    for sc in examples:
        pm = getattr(sc, "preds_missing", [])
        tm = getattr(sc, "trans_missing", [])
        bits.append(f"- `{sc.name}` missing preds={pm or '—'} trans={tm or '—'}")

    return [
        Recommendation(
            text=(
                "Some scenarios reference predicates/transitions that do not exist yet.\n"
                + "\n".join(bits)
                + ("\n(…more)" if len(missing) > len(examples) else "")
                + "\nRecommendation: add the missing predicates/transitions or edit the scenarios to match the domain model."
            )
        )
    ]


def _rule_inconsistent_scenarios(
    old_inst: StateInstance,
    new_inst: StateInstance,
    failure: ChangeFailure | None,
) -> Iterable[Recommendation]:
    if not _ok(failure):
        return []

    spec = new_inst.spec
    bad = [sc for sc in spec.scenarios if getattr(sc, "all_preds_consistent", True) is False]
    if not bad:
        return []

    names = ", ".join(f"`{sc.name}`" for sc in bad[:5])
    tail = " …" if len(bad) > 5 else ""
    return [
        Recommendation(
            text=(
                f"Some scenarios appear **inconsistent** (unsat) given the current predicates: {names}{tail}.\n"
                "Recommendation: relax contradictory predicates, or split the scenario into smaller cases."
            )
        )
    ]


def _rule_conflicting_scenarios(
    old_inst: StateInstance,
    new_inst: StateInstance,
    failure: ChangeFailure | None,
) -> Iterable[Recommendation]:
    if not _ok(failure):
        return []

    spec = new_inst.spec
    conflicted = [sc for sc in spec.scenarios if getattr(sc, "conflicts", [])]
    if not conflicted:
        return []

    names = ", ".join(f"`{sc.name}`" for sc in conflicted[:5])
    tail = " …" if len(conflicted) > 5 else ""
    return [
        Recommendation(
            text=(
                f"Some scenarios have **overlaps/conflicts** with others: {names}{tail}.\n"
                "Recommendation: add distinguishing predicates (boundary conditions) so scenarios become disjoint, "
                "or explicitly document intended priority/ordering."
            )
        )
    ]


def _rule_unmatched_traces(
    old_inst: StateInstance,
    new_inst: StateInstance,
    failure: ChangeFailure | None,
) -> Iterable[Recommendation]:
    if not _ok(failure):
        return []

    cont = new_inst.art_container
    amap = new_inst.art_map

    if amap is None:
        return []

    def _is_trace_matched(trace) -> bool:
        tid = getattr(trace, "art_id", None)
        if not tid:
            return False
        return bool(amap.get_comp_for_artifact(tid))

    all_traces = list(cont.test_traces) + list(cont.log_traces)
    if not all_traces:
        return []

    unmatched = [t for t in all_traces if not _is_trace_matched(t)]
    if not unmatched:
        return []

    def _label(t) -> str:
        return (
            getattr(t, "name", None)
            or getattr(t, "filename", None)
            or (getattr(t, "art_id", None) or "<unknown>")
        )

    examples = unmatched[:3]
    bullets = "\n".join(f"- `{_label(t)}`" for t in examples)
    return [
        Recommendation(
            text=(
                "Some traces are not matched to any scenario.\n"
                f"{bullets}"
                + ("\n(…more)" if len(unmatched) > len(examples) else "")
                + "\nRecommendation: add/adjust scenarios so each important trace is explained "
                "(or explicitly link traces to components if you support manual traceability)."
            )
        )
    ]


def _rule_invalid_iml_components(
    old_inst: StateInstance,
    new_inst: StateInstance,
    failure: ChangeFailure | None,
) -> Iterable[Recommendation]:
    if not _ok(failure):
        return []

    dm = new_inst.spec.domain_model

    bad_preds = []
    for p in list(getattr(dm, "state_preds", [])) + list(getattr(dm, "action_preds", [])):
        if getattr(p, "is_iml_valid", None) and str(getattr(p, "is_iml_valid")) == "IMLValidity.INVALID":
            bad_preds.append(p)

    bad_trans = [
        t
        for t in getattr(dm, "transitions", [])
        if getattr(t, "is_iml_valid", None) and str(getattr(t, "is_iml_valid")) == "IMLValidity.INVALID"
    ]

    if not bad_preds and not bad_trans:
        return []

    lines: list[str] = []
    if bad_preds:
        lines.append(
            f"- Invalid predicates: {', '.join(f'`{p.name}`' for p in bad_preds[:6])}{' …' if len(bad_preds) > 6 else ''}"
        )
    if bad_trans:
        lines.append(
            f"- Invalid transitions: {', '.join(f'`{t.name}`' for t in bad_trans[:6])}{' …' if len(bad_trans) > 6 else ''}"
        )

    return [
        Recommendation(
            kind="warning",
            text=(
                "Some model components fail IML/typechecking.\n"
                + "\n".join(lines)
                + "\nRecommendation: fix IML errors first; scenario matching and trace checking won’t be trustworthy otherwise."
            ),
        )
    ]


def _rule_traceability_nudge(
    old_inst: StateInstance,
    new_inst: StateInstance,
    failure: ChangeFailure | None,
) -> Iterable[Recommendation]:
    if not _ok(failure):
        return []

    amap = new_inst.art_map
    if amap is None:
        return []

    n_arts = (
        len(new_inst.art_container.test_traces)
        + len(new_inst.art_container.log_traces)
        + len(new_inst.art_container.doc_ref)
        + len(new_inst.art_container.src_code)
    )
    n_links = sum(len(v) for v in amap.art_to_comp_map.values())

    if n_arts >= 3 and n_links == 0:
        return [
            Recommendation(
                text=(
                    "You have artifacts loaded, but **no traceability links** yet (ArtifactMap is empty).\n"
                    "Recommendation: run trace matching (or add explicit links) so docs/src/traces connect to "
                    "predicates/transitions/scenarios. This makes the synthesized spec explainable."
                )
            )
        ]

    return []


def _rule_state_diff_summary(
    old_inst: StateInstance,
    new_inst: StateInstance,
    failure: ChangeFailure | None,
) -> Iterable[Recommendation]:
    if not _ok(failure):
        return []

    sd: StateDiff | None = new_inst.state_diff
    if sd is None:
        return []

    recs: list[Recommendation] = []

    # 1) If base got better/worse, say that first
    base_out = _outcome(sd.base_status)  # type: ignore[arg-type]
    if base_out in (ComparisonOutcome.IMPROVED, ComparisonOutcome.DECLINED):
        if base_out is ComparisonOutcome.IMPROVED:
            recs.append(
                _rec(
                    f"Domain model base improved ({sd.base_status.before} → {sd.base_status.after}). "
                    "Proceed to checking predicates/transitions and scenario consistency.",
                    kind="info",
                    priority=10,
                )
            )
        else:
            recs.append(
                _rec(
                    f"Domain model base declined ({sd.base_status.before} → {sd.base_status.after}). "
                    "Fix base IML/type issues before trusting any further analysis.",
                    kind="error",
                    priority=0,
                )
            )

    # 2) Predicates
    if _outcome(sd.num_preds_errored) is ComparisonOutcome.IMPROVED:  # type: ignore[arg-type]
        recs.append(
            _rec(
                f"Fewer predicate errors ({sd.num_preds_errored.before} → {sd.num_preds_errored.after}). "
                "Re-run scenario consistency + trace matching to leverage the improved predicate set.",
                kind="info",
                priority=20,
            )
        )
    elif _outcome(sd.num_preds_errored) is ComparisonOutcome.DECLINED:  # type: ignore[arg-type]
        recs.append(
            _rec(
                f"More predicate errors ({sd.num_preds_errored.before} → {sd.num_preds_errored.after}). "
                "Inspect the most recently added/edited predicates and fix their IML first.",
                kind="error",
                priority=5,
            )
        )

    st_err_out = _outcome(sd.num_state_preds_errored)  # type: ignore[arg-type]
    if st_err_out in (ComparisonOutcome.IMPROVED, ComparisonOutcome.DECLINED):
        recs.append(
            _rec(
                f"State-predicate errors changed ({sd.num_state_preds_errored.before} → {sd.num_state_preds_errored.after}). "
                "Focus on `state -> bool` predicates first; they gate most scenario `given` checks.",
                kind="info" if st_err_out is ComparisonOutcome.IMPROVED else "warning",
                priority=23 if st_err_out is ComparisonOutcome.IMPROVED else 9,
            )
        )

    act_err_out = _outcome(sd.num_action_preds_errored)  # type: ignore[arg-type]
    if act_err_out in (ComparisonOutcome.IMPROVED, ComparisonOutcome.DECLINED):
        recs.append(
            _rec(
                f"Action-predicate errors changed ({sd.num_action_preds_errored.before} → {sd.num_action_preds_errored.after}). "
                "Fix `state * action -> bool` predicates; they affect scenario `when` and trace matching.",
                kind="info" if act_err_out is ComparisonOutcome.IMPROVED else "warning",
                priority=24 if act_err_out is ComparisonOutcome.IMPROVED else 10,
            )
        )

    if _outcome(sd.num_preds_matched) is ComparisonOutcome.IMPROVED:  # type: ignore[arg-type]
        recs.append(
            _rec(
                f"More predicates are now matched to artifacts ({sd.num_preds_matched.before} → {sd.num_preds_matched.after}). "
                "Explain these matches in prose and add missing traceability for remaining artifacts.",
                kind="next",
                priority=30,
            )
        )

    # 3) Transitions
    if _outcome(sd.num_trans_valid_logic) is ComparisonOutcome.IMPROVED:  # type: ignore[arg-type]
        recs.append(
            _rec(
                f"More transitions are logic-valid ({sd.num_trans_valid_logic.before} → {sd.num_trans_valid_logic.after}). "
                "Attach these transitions to scenarios’ `then` lists so traces can match end-to-end.",
                kind="next",
                priority=25,
            )
        )
    elif _outcome(sd.num_trans_valid_logic) is ComparisonOutcome.DECLINED:  # type: ignore[arg-type]
        recs.append(
            _rec(
                f"Fewer transitions are logic-valid ({sd.num_trans_valid_logic.before} → {sd.num_trans_valid_logic.after}). "
                "Revert or fix the last transition edit; invalid transitions break scenario execution and matching.",
                kind="error",
                priority=5,
            )
        )

    # 4) Scenarios
    if _outcome(sd.num_sc_missing) is ComparisonOutcome.DECLINED:
        recs.append(
            _rec(
                f"More scenarios are missing components ({sd.num_sc_missing.before} → {sd.num_sc_missing.after}). "
                "Add the referenced predicates/transitions or remove them from those scenarios.",
                kind="error",
                priority=8,
            )
        )
    elif _outcome(sd.num_sc_missing) is ComparisonOutcome.IMPROVED:
        recs.append(
            _rec(
                f"Fewer scenarios are missing components ({sd.num_sc_missing.before} → {sd.num_sc_missing.after}). "
                "Re-check scenario consistency and conflicts now that they are fully specified.",
                kind="next",
                priority=18,
            )
        )

    if _outcome(sd.num_sc_inconsistent) is ComparisonOutcome.DECLINED:
        recs.append(
            _rec(
                f"More scenarios are inconsistent ({sd.num_sc_inconsistent.before} → {sd.num_sc_inconsistent.after}). "
                "Split scenarios or relax contradictory predicate combinations.",
                kind="error",
                priority=12,
            )
        )
    elif _outcome(sd.num_sc_inconsistent) is ComparisonOutcome.IMPROVED:
        recs.append(
            _rec(
                f"Fewer scenarios are inconsistent ({sd.num_sc_inconsistent.before} → {sd.num_sc_inconsistent.after}). "
                "You’re in a good place to match traces and generate human-readable explanations.",
                kind="info",
                priority=22,
            )
        )

    if _outcome(sd.num_sc_conflicted) is ComparisonOutcome.DECLINED:
        recs.append(
            _rec(
                f"More scenario conflicts/overlaps detected ({sd.num_sc_conflicted.before} → {sd.num_sc_conflicted.after}). "
                "Add boundary predicates to make scenarios disjoint, or document intended priority.",
                kind="error",
                priority=15,
            )
        )
    elif _outcome(sd.num_sc_conflicted) is ComparisonOutcome.IMPROVED:
        recs.append(
            _rec(
                f"Fewer scenario conflicts/overlaps ({sd.num_sc_conflicted.before} → {sd.num_sc_conflicted.after}). "
                "Proceed to coverage work—ensure each important trace has at least one matching scenario.",
                kind="next",
                priority=28,
            )
        )

    # 5) Trace coverage
    if _outcome(sd.num_test_traces_logic_good) is ComparisonOutcome.DECLINED:
        recs.append(
            _rec(
                f"More test traces are failing IML validity ({sd.num_test_traces_logic_good.before} → {sd.num_test_traces_logic_good.after}). "
                "Fix the trace encodings (given/when/then) so matching is meaningful.",
                kind="error",
                priority=10,
            )
        )

    if _outcome(sd.num_test_traces_matched) is ComparisonOutcome.IMPROVED:
        recs.append(
            _rec(
                f"More test traces matched scenarios ({sd.num_test_traces_matched.before} → {sd.num_test_traces_matched.after}). "
                "Summarize the matched scenarios as the system’s inferred behavior for those tests.",
                kind="next",
                priority=35,
            )
        )

    if _outcome(sd.num_log_traces_matched) is ComparisonOutcome.IMPROVED:
        recs.append(
            _rec(
                f"More log traces matched scenarios ({sd.num_log_traces_matched.before} → {sd.num_log_traces_matched.after}). "
                "Treat these as in-the-wild confirmation and use them to justify model boundaries.",
                kind="next",
                priority=40,
            )
        )

    if not recs:
        recs.append(
            _rec(
                "No major metric movement detected in the last change. "
                "Add a predicate/transition/scenario (or new traces) to drive measurable progress.",
                kind="info",
                priority=60,
            )
        )

    return recs