#
#   Imandra Inc.
#
#   speclogician/presentation/builders/trace_reachability.py
#

from __future__ import annotations

from dataclasses import dataclass

from speclogician.data.container import ArtifactContainer
from speclogician.modeling.trace_reachability import TraceReachabilityAnalysis
from speclogician.presentation.ctx import RenderCtx
from speclogician.presentation.models.trace_reachability import (
    TraceReachabilityPM,
    TraceReachabilityCountsPM,
    TraceReachabilitySummaryPM,
    TraceReachabilityTracePM,
    TraceReachabilityHopPM,
)


@dataclass(frozen=True)
class TraceReachabilityBuildOpts:
    max_traces: int = 50
    max_unhit_preview: int = 20
    max_unsupported_hops_per_trace: int = 10


def _trace_label(cont: ArtifactContainer, trace_id: str) -> tuple[str, str]:
    for t in cont.test_traces:
        if t.art_id == trace_id:
            return "test", (t.name or trace_id)
    for t in cont.log_traces:
        if t.art_id == trace_id:
            return "log", (t.filename or trace_id)
    return "unknown", trace_id


def build_trace_reachability_pm(
    ana: TraceReachabilityAnalysis | None,
    art_cont: ArtifactContainer,
    ctx: RenderCtx,
    *,
    opts: TraceReachabilityBuildOpts = TraceReachabilityBuildOpts(),
) -> TraceReachabilityPM | None:
    if ana is None:
        return None

    s = ana.summary

    pm = TraceReachabilityPM(
        spec_fp=getattr(ana, "spec_fp", None),
        artifacts_fp=getattr(ana, "artifacts_fp", None),
        mapping_fp=getattr(ana, "mapping_fp", None),
        counts=TraceReachabilityCountsPM(
            total_traces=s.total_traces,
            total_steps=s.total_steps,
            total_hops=s.total_hops,
            supported_hops=s.supported_hops,
            unsupported_hops=s.unsupported_hops,
            scenarios_hit=len(list(s.scenarios_hit or [])),
            scenarios_unhit=len(list(s.scenarios_unhit or [])),
        ),
        summary=TraceReachabilitySummaryPM(
            top_unsupported_edges=list(s.top_unsupported_edges or []),
            top_frontier_trace_endings=list(s.top_frontier_trace_endings or []),
            scenarios_unhit_preview=list(s.scenarios_unhit or [])[: opts.max_unhit_preview],
        ),
    )

    # ------------------------------------------------------------
    # Filter: if ctx.trace_id is set, only include that trace.
    # Unknown trace_id => empty list (counts/summary still global).
    # ------------------------------------------------------------
    wanted = getattr(ctx, "trace_id", None)
    if wanted:
        per_trace_items = [(wanted, (ana.per_trace or {}).get(wanted))]
        per_trace_items = [(tid, al) for (tid, al) in per_trace_items if al is not None]
    else:
        per_trace_items = list((ana.per_trace or {}).items())

    rows: list[TraceReachabilityTracePM] = []

    for trace_id, al in per_trace_items:
        kind, label = _trace_label(art_cont, trace_id)

        # collect unsupported hops (capped)
        uns: list[TraceReachabilityHopPM] = []
        for h in (al.observed_hops or []):
            if h.supported_by_graph:
                continue
            uns.append(
                TraceReachabilityHopPM(
                    step_from=h.step_from,
                    step_to=h.step_to,
                    src=h.src if h.src else None,
                    dst=h.dst if h.dst else None,
                    supported_by_graph=bool(h.supported_by_graph),
                    reason=h.reason,
                    shortest_path=h.shortest_path,
                )
            )
            if len(uns) >= opts.max_unsupported_hops_per_trace:
                break

        rows.append(
            TraceReachabilityTracePM(
                trace_id=trace_id,
                trace_label=label,
                trace_kind=kind,
                num_steps=al.num_steps,
                num_matched_steps=al.num_matched_steps,
                num_supported_hops=al.num_supported_hops,
                num_unsupported_hops=al.num_unsupported_hops,
                step_best=list(al.step_best or []),
                unsupported_hops=uns,
            )
        )

    # sort “worst first” for list view
    rows.sort(
        key=lambda r: (
            -r.num_unsupported_hops,
            -r.num_steps,
            r.trace_label.lower(),
            r.trace_id,
        )
    )

    pm.traces = rows[: opts.max_traces]
    return pm