#
#   Imandra Inc.
#   
#   speclogician/logic/lib/transitions.py
#

import asyncio

from rich.progress import Progress

from speclogician.modeling.spec import Spec
from speclogician.logic.api.client import SpecLogicianImandraX
from speclogician.utils import IMLValidity, console

def check_transitions(
    spec: Spec,
    names: list[str] = [],
    json_only: bool = False,
) -> None:
    """
    Check transition IML validity via ImandraX and write results into transition objects:
        - transition.is_iml_valid
        - transition.iml_error
    If `names` is non-empty, checks only those transition names; otherwise checks all.
    """

    dm = spec.domain_model
    all_trans = list(dm.transitions)

    if names:
        name_set = set(names)
        trans_to_check = [t for t in all_trans if t.name in name_set]

        if not json_only:
            known = {t.name for t in all_trans}
            missing = [n for n in names if n not in known]
            if missing:
                console.print(f"[yellow]Logician[/yellow] Unknown transition name(s) requested: {missing}")
    else:
        trans_to_check = all_trans

    def _errs_to_str(res: object) -> str:
        es = getattr(res, "errors", None) or []
        if es:
            return "\n".join(str(e) for e in es)
        return "ImandraX eval_model: has_errors=True"

    if not json_only:
        console.print("[bold cyan]Logician[/bold cyan] Checking transition IML validity")

    async def _check_all() -> None:
        async with SpecLogicianImandraX() as ix:
            for t in trans_to_check:
                src = dm.make_iml_model([], [], [t.name])
                try:
                    r = await ix.load_model(src)
                    has_err = bool(getattr(r, "has_errors", False))
                    t.is_iml_valid = IMLValidity.INVALID if has_err else IMLValidity.VALID
                    t.iml_error = None if not has_err else _errs_to_str(r)
                except Exception as e:
                    t.is_iml_valid = IMLValidity.UNKNOWN
                    t.iml_error = str(e)

    if not json_only:
        with Progress(transient=True) as progress:
            task = progress.add_task("Checking transitions", total=len(trans_to_check))

            async def _check_with_progress() -> None:
                async with SpecLogicianImandraX() as ix:
                    for t in trans_to_check:
                        src = dm.make_iml_model([], [], [t.name])
                        try:
                            r = await ix.load_model(src)
                            has_err = bool(getattr(r, "has_errors", False))
                            t.is_iml_valid = IMLValidity.INVALID if has_err else IMLValidity.VALID
                            t.iml_error = None if not has_err else _errs_to_str(r)
                        except Exception as e:
                            t.is_iml_valid = IMLValidity.UNKNOWN
                            t.iml_error = str(e)
                        progress.advance(task)

            try:
                asyncio.run(_check_with_progress())
            except Exception as e:
                for t in trans_to_check:
                    t.is_iml_valid = IMLValidity.UNKNOWN
                    t.iml_error = str(e)
    else:
        try:
            asyncio.run(_check_all())
        except Exception as e:
            for t in trans_to_check:
                t.is_iml_valid = IMLValidity.UNKNOWN
                t.iml_error = str(e)

    # Refresh domain-model transition stats (global, not just filtered set)
    dm.num_trans_total = len(dm.transitions)
    dm.num_trans_valid_logic = sum(1 for t in dm.transitions if t.is_iml_valid == IMLValidity.VALID)

    if not json_only:
        console.print(
            f"[green]✓[/green] Transitions checked. "
            f"valid={dm.num_trans_valid_logic}/{dm.num_trans_total}"
        )
