#
#   Imandra Inc.
#
#   speclogician/modeling/complement.py
#

from __future__ import annotations

import hashlib
import json
from typing import Sequence, Iterable

from pydantic import BaseModel
from imandrax_api_models import RegionStr


class ScenarioComplement(BaseModel):
    """Result of running region decomposition on all the combined scenarios."""
    regions: Sequence[RegionStr]


class ScenarioComplementDiff(BaseModel):
    """
    Difference between two scenario complements (before -> after),
    computed using stable fingerprints of regions.

    We store only fingerprints by default (regions can be huge).
    """
    num_regions_before: int
    num_regions_after: int
    regions_added: list[str]      # fingerprints that appear in after but not before
    regions_removed: list[str]    # fingerprints that appear in before but not after


def _canon_region_str(r: RegionStr) -> str:
    """
    Produce a deterministic string representation of a RegionStr.

    We prefer a canonical JSON-ish blob rather than str(r) because:
      - __repr__ / __str__ may change
      - dict ordering may vary
    """
    # constraints_str is typically a list of strings
    constraints = list(getattr(r, "constraints_str", []) or [])

    invariant = getattr(r, "invariant_str", "") or ""
    model_eval = getattr(r, "model_eval_str", "") or ""

    # model_str can be dict-like (often {'s': '{x = 0}'}). Canonicalize ordering.
    model = getattr(r, "model_str", None)
    if isinstance(model, dict):
        # stringify values to avoid non-jsonable shapes
        model_items = {str(k): str(v) for k, v in sorted(model.items(), key=lambda kv: str(kv[0]))}
    elif model is None:
        model_items = {}
    else:
        # fallback: stringify
        model_items = {"_": str(model)}

    payload = {
        "constraints": constraints,
        "invariant": invariant,
        "model": model_items,
        "model_eval": model_eval,
    }

    # separators for stable compact encoding
    return json.dumps(payload, sort_keys=True, separators=(",", ":"))


def _fingerprint_region(r: RegionStr | str) -> str:
    """
    Stable-ish fingerprint for a RegionStr (or a raw string region).
    """
    if isinstance(r, RegionStr):
        s = _canon_region_str(r)
    else:
        s = str(r)

    h = hashlib.sha256(s.encode("utf-8")).hexdigest()
    # short prefix is plenty for diffs + UI, but keep full if you prefer
    return h[:16]


def _fingerprints(regions: Iterable[RegionStr | str]) -> set[str]:
    return {_fingerprint_region(r) for r in regions}


def diff_scenario_complement(before: ScenarioComplement, after: ScenarioComplement) -> ScenarioComplementDiff:
    before_regions = list(before.regions or [])
    after_regions = list(after.regions or [])

    bf = _fingerprints(before_regions)
    af = _fingerprints(after_regions)

    added = sorted(list(af - bf))
    removed = sorted(list(bf - af))

    return ScenarioComplementDiff(
        num_regions_before=len(before_regions),
        num_regions_after=len(after_regions),
        regions_added=added,
        regions_removed=removed,
    )