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

import asyncio

from rich.progress import Progress, TaskID

from speclogician.data.container import ArtifactContainer
from speclogician.data.mapping import ArtifactMap, ElementType
from speclogician.data.traces import TraceArtifact, TraceScenarioMatchResult, TraceIMLValidity
from speclogician.logic.api.client import SpecLogicianImandraX
from speclogician.modeling.spec import Spec
from speclogician.utils import IMLValidity, console

def check_traces(
    spec: Spec,
    art_cont: ArtifactContainer,
    art_map: ArtifactMap,
    traces: list[str],
    json_only: bool = False,
) -> None:
    dm = spec.domain_model

    all_traces: list[TraceArtifact] = []
    all_traces.extend(list(art_cont.test_traces))
    all_traces.extend(list(art_cont.log_traces))

    if traces:
        want = set(traces)
        to_check = [t for t in all_traces if t.art_id in want]
        if not json_only:
            missing = [tid for tid in traces if tid not in {t.art_id for t in all_traces}]
            if missing:
                console.print(f"[yellow]Logician[/yellow] Unknown trace id(s) requested: {missing}")
    else:
        to_check = all_traces

    if not to_check:
        if not json_only:
            console.print("[bold cyan]Logician[/bold cyan] No traces to check.")
        return

    def _scenario_match_ok(res : TraceScenarioMatchResult) -> bool:
        # TraceScenarioMatchResult has bool flags; mismatch is not an error
        return bool(res.iml_valid) and res.given_matches and res.when_matches and res.transition_matches

    async def _run_checks(progress: Progress | None = None, task_trace:TaskID|None=None) -> None:
        async with SpecLogicianImandraX() as ix:
            for tr in to_check:

                if tr.art_id is None: # This should never be the case as we check this on the pydantic level
                    continue

                # Clear old links for this artifact to avoid stale mapping
                art_map.rem_element(ElementType.ArtifactType, tr.art_id)

                # 1) IML validity
                try:
                    v = await ix.check_trace_iml_validity(tr, dm)
                    # v: TraceIMLValidity(iml_valid: IMLValidity, err: str|None)
                    tr.iml_validity = v
                except Exception as e:
                    tr.iml_validity = TraceIMLValidity(iml_valid=IMLValidity.UNKNOWN, err=str(e))
                    if progress is not None and task_trace is not None:
                        progress.advance(task_trace)
                    continue

                # Only proceed to scenario matching if trace is VALID IML
                if tr.iml_validity.iml_valid != IMLValidity.VALID:
                    if progress is not None and task_trace is not None:
                        progress.advance(task_trace)
                    continue

                # 2) Match against scenarios; record matches
                for sc in spec.scenarios:
                    try:
                        m : TraceScenarioMatchResult = await ix.check_trace_match(tr, sc, dm)
                    except Exception:
                        # tool failure => skip (don't manufacture mismatch errors)
                        continue

                    if _scenario_match_ok(m):
                        art_map.add_connection(art_id=tr.art_id, comp_name=sc.comp_id)

                if progress is not None and task_trace is not None:
                    progress.advance(task_trace)

    if not json_only:
        console.print(
            f"[bold cyan]Logician[/bold cyan] Checking traces ({len(to_check)}) "
            f"against scenarios ({len(spec.scenarios)})"
        )

        with Progress(transient=True) as progress:
            task_trace = progress.add_task("Checking traces", total=len(to_check))
            try:
                asyncio.run(_run_checks(progress=progress, task_trace=task_trace))
            except Exception as e:
                console.print(f"[red]Logician[/red] Trace checking failed: {e}")
                return

        n_valid = sum(1 for t in to_check if t.iml_validity.iml_valid == IMLValidity.VALID)
        n_invalid = sum(1 for t in to_check if t.iml_validity.iml_valid == IMLValidity.INVALID)
        n_unknown = sum(1 for t in to_check if t.iml_validity.iml_valid == IMLValidity.UNKNOWN)
        console.print(f"[green]✓[/green] Traces checked. valid={n_valid}, invalid={n_invalid}, unknown={n_unknown}")
    else:
        try:
            asyncio.run(_run_checks())
        except Exception:
            return
    