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

import asyncio

from rich.progress import Progress

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

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

    dm = spec.domain_model

    # Resolve which predicates to check (by name), across state+action predicates.
    all_state = list(dm.state_preds)
    all_action = list(dm.action_preds)

    if names:
        name_set = set(names)
        state_to_check = [p for p in all_state if p.name in name_set]
        action_to_check = [p for p in all_action if p.name in name_set]

        # If a name was requested but doesn't exist, we can't attach the result to a predicate.
        # We *log* it (non-json) but do not raise.
        if not json_only:
            known = {p.name for p in all_state} | {p.name for p in all_action}
            missing = [n for n in names if n not in known]
            if missing:
                console.print(f"[yellow]Logician[/yellow] Unknown predicate name(s) requested: {missing}")
    else:
        state_to_check = all_state
        action_to_check = all_action

    # Helper: stringify ImandraX eval_model errors (best effort)
    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"

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

            # --- Action predicates ---
            for p in action_to_check:
                src = dm.make_iml_model([], [p.name], [])
                try:
                    r = await ix.load_model(src)
                    has_err = bool(getattr(r, "has_errors", False))
                    p.is_iml_valid = IMLValidity.INVALID if has_err else IMLValidity.VALID
                    p.iml_error = None if not has_err else _errs_to_str(r)
                except Exception as e:
                    p.is_iml_valid = IMLValidity.UNKNOWN
                    p.iml_error = str(e)

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

    # Show progress only in non-json mode
    if not json_only:
        with Progress(transient=True) as progress:
            t_state = progress.add_task("Checking state predicates", total=len(state_to_check))
            t_action = progress.add_task("Checking action predicates", total=len(action_to_check))

            # We still do the actual work in one ImandraX session; progress is updated as we go.
            async def _check_with_progress() -> None:
                async with SpecLogicianImandraX() as ix:
                    for p in state_to_check:
                        src = dm.make_iml_model([p.name], [], [])
                        try:
                            r = await ix.load_model(src)
                            has_err = bool(getattr(r, "has_errors", False))
                            p.is_iml_valid = IMLValidity.INVALID if has_err else IMLValidity.VALID
                            p.iml_error = None if not has_err else _errs_to_str(r)
                        except Exception as e:
                            p.is_iml_valid = IMLValidity.UNKNOWN
                            p.iml_error = str(e)
                        progress.advance(t_state)

                    for p in action_to_check:
                        src = dm.make_iml_model([], [p.name], [])
                        try:
                            r = await ix.load_model(src)
                            has_err = bool(getattr(r, "has_errors", False))
                            p.is_iml_valid = IMLValidity.INVALID if has_err else IMLValidity.VALID
                            p.iml_error = None if not has_err else _errs_to_str(r)
                        except Exception as e:
                            p.is_iml_valid = IMLValidity.UNKNOWN
                            p.iml_error = str(e)
                        progress.advance(t_action)

            try:
                asyncio.run(_check_with_progress())
            except Exception as e:
                # If the whole checking run fails (e.g. credentials/network),
                # we don't have per-predicate errors; mark checked ones as unknown.
                for p in state_to_check + action_to_check:
                    p.is_iml_valid = IMLValidity.UNKNOWN
                    p.iml_error = str(e)
    else:
        # No progress bar, just run.
        try:
            asyncio.run(_check_all())
        except Exception as e:
            for p in state_to_check + action_to_check:
                p.is_iml_valid = IMLValidity.UNKNOWN
                p.iml_error = str(e)

    # Refresh domain-model predicate stats (global, not just filtered set)
    dm.num_state_preds_total = len(all_state)
    dm.num_action_preds_total = len(all_action)

    dm.num_state_preds_valid_logic = sum(1 for p in all_state if getattr(p, "is_iml_valid", IMLValidity.UNKNOWN) == IMLValidity.VALID)
    dm.num_action_preds_valid_logic = sum(1 for p in all_action if getattr(p, "is_iml_valid", IMLValidity.UNKNOWN) == IMLValidity.VALID)

    dm.num_preds_total = dm.num_state_preds_total + dm.num_action_preds_total
    dm.num_preds_valid_logic = dm.num_state_preds_valid_logic + dm.num_action_preds_valid_logic

    if not json_only:
        console.print(
            f"[green]✓[/green] Predicates checked. "
            f"state_valid={dm.num_state_preds_valid_logic}/{dm.num_state_preds_total}, "
            f"action_valid={dm.num_action_preds_valid_logic}/{dm.num_action_preds_total}"
        )
    
