"""Offline Hugging Face Space for Trace2MCP.

The UI is intentionally self-contained: it never calls a model or a remote API.
It can use the Trace2MCP package when installed and falls back to a transparent
graph preview while the package is being developed.
"""

from __future__ import annotations

import html
import json
import os
import sys
from dataclasses import asdict, is_dataclass
from datetime import datetime
from pathlib import Path
from typing import Any

# Keep the public demo genuinely offline, including framework analytics and
# version checks. The compiler itself never needs a model, token, or network.
os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "False")
os.environ.setdefault("HF_HUB_DISABLE_TELEMETRY", "1")

import gradio as gr

ROOT = Path(__file__).resolve().parent
SAMPLE_PATH = ROOT / "examples" / "sample_trace.json"
STYLE_PATH = ROOT / "assets" / "style.css"
SRC_PATH = ROOT / "src"
if SRC_PATH.exists() and str(SRC_PATH) not in sys.path:
    sys.path.insert(0, str(SRC_PATH))


def _load_engine() -> tuple[Any | None, Any | None, str]:
    """Load the public API without making the Space depend on one import layout."""

    try:
        from trace2mcp.compiler import compile_trace
        from trace2mcp.parsers import load_trace

        return load_trace, compile_trace, "package"
    except Exception as first_error:  # pragma: no cover - depends on install state
        try:
            from trace2mcp import compile_trace, load_trace

            return load_trace, compile_trace, "package"
        except Exception:
            return None, None, f"preview ({type(first_error).__name__})"


LOAD_TRACE, COMPILE_TRACE, ENGINE_MODE = _load_engine()


def _parse_with_package(raw: dict[str, Any]) -> Any:
    """Parse an in-memory Trace v1 using either supported package surface."""

    if LOAD_TRACE is None:
        return raw
    try:
        return LOAD_TRACE(raw)
    except TypeError:
        # The alpha loader accepts filesystem paths; the model is the canonical
        # in-memory equivalent used by the interactive editor.
        from trace2mcp.models import Trace

        return Trace.from_dict(raw)


FALLBACK_SAMPLE: dict[str, Any] = {
    "version": "1",
    "name": "company_brief",
    "inputs": {"company": "Acme Labs"},
    "events": [
        {
            "id": "resolve",
            "tool": "resolve_company",
            "duration_ms": 520,
            "depends_on": [],
            "side_effect": "read",
            "arguments": {"company": "Acme Labs"},
        },
        {
            "id": "news",
            "tool": "search_news",
            "duration_ms": 900,
            "depends_on": ["resolve"],
            "side_effect": "read",
            "arguments": {"company_id": "acme-labs"},
        },
        {
            "id": "profile",
            "tool": "get_company_profile",
            "duration_ms": 730,
            "depends_on": ["resolve"],
            "side_effect": "read",
            "arguments": {"company_id": "acme-labs"},
        },
        {
            "id": "financials",
            "tool": "get_financial_snapshot",
            "duration_ms": 840,
            "depends_on": ["resolve"],
            "side_effect": "read",
            "arguments": {"company_id": "acme-labs"},
        },
        {
            "id": "rank",
            "tool": "rank_evidence",
            "duration_ms": 310,
            "depends_on": ["news", "profile", "financials"],
            "side_effect": "read",
            "arguments": {
                "news": "{{ news.items }}",
                "profile": "{{ profile }}",
                "financials": "{{ financials }}",
            },
        },
        {
            "id": "render",
            "tool": "render_brief",
            "duration_ms": 180,
            "depends_on": ["rank"],
            "side_effect": "read",
            "arguments": {"evidence": ["Verified company profile", "Recent public roadmap"]},
        },
    ],
}


def _sample_text() -> str:
    if SAMPLE_PATH.exists():
        return SAMPLE_PATH.read_text(encoding="utf-8")
    return json.dumps(FALLBACK_SAMPLE, indent=2, ensure_ascii=False)


def _trace_mapping(value: Any) -> dict[str, Any]:
    if isinstance(value, dict):
        return value
    for method_name in ("model_dump", "to_dict", "dict"):
        method = getattr(value, method_name, None)
        if callable(method):
            result = method()
            if isinstance(result, dict):
                return result
    if is_dataclass(value):
        return asdict(value)
    if hasattr(value, "__dict__"):
        return dict(vars(value))
    raise TypeError("The trace loader did not return a mapping-like object.")


def _serializable(value: Any) -> Any:
    if value is None or isinstance(value, (str, int, float, bool)):
        return value
    if isinstance(value, dict):
        return {str(key): _serializable(item) for key, item in value.items()}
    if isinstance(value, (list, tuple, set)):
        return [_serializable(item) for item in value]
    for method_name in ("model_dump", "to_dict", "dict"):
        method = getattr(value, method_name, None)
        if callable(method):
            return _serializable(method())
    if is_dataclass(value):
        return _serializable(asdict(value))
    if hasattr(value, "__dict__"):
        return _serializable(vars(value))
    return repr(value)


def _calls_from(trace: dict[str, Any]) -> list[dict[str, Any]]:
    raw_calls = trace.get("calls") or trace.get("steps") or trace.get("events")
    if not isinstance(raw_calls, list) or not raw_calls:
        raise ValueError("Expected a non-empty 'calls', 'steps', or 'events' array.")

    calls: list[dict[str, Any]] = []
    for index, raw in enumerate(raw_calls):
        if not isinstance(raw, dict):
            raise ValueError(f"Step {index + 1} must be a JSON object.")
        call_id = str(raw.get("id") or raw.get("call_id") or f"step_{index + 1}")
        tool = str(raw.get("tool") or raw.get("name") or raw.get("tool_name") or "tool")
        duration = raw.get("duration_ms", raw.get("latency_ms"))
        started_at = raw.get("started_at")
        ended_at = raw.get("ended_at")
        if started_at and ended_at:
            try:
                started = datetime.fromisoformat(str(started_at).replace("Z", "+00:00"))
                ended = datetime.fromisoformat(str(ended_at).replace("Z", "+00:00"))
                duration = (ended - started).total_seconds() * 1000
            except ValueError as error:
                raise ValueError(f"'{call_id}' has invalid ISO timestamps.") from error
        if duration is None:
            duration = 0
        try:
            duration_ms = max(0.0, float(duration))
        except (TypeError, ValueError) as error:
            raise ValueError(f"'{call_id}' has an invalid duration_ms value.") from error
        dependencies = raw.get("depends_on") or raw.get("dependencies") or []
        if isinstance(dependencies, str):
            dependencies = [dependencies]
        if not isinstance(dependencies, list):
            raise ValueError(f"'{call_id}' depends_on must be an array.")
        if "read_only" in raw:
            read_only = bool(raw["read_only"])
        else:
            read_only = raw.get("side_effect", "read") == "read"
        calls.append(
            {
                **raw,
                "id": call_id,
                "tool": tool,
                "duration_ms": duration_ms,
                "depends_on": [str(item) for item in dependencies],
                "read_only": read_only,
            }
        )

    ids = [call["id"] for call in calls]
    if len(ids) != len(set(ids)):
        raise ValueError("Every step must have a unique id.")
    known = set(ids)
    for call in calls:
        missing = set(call["depends_on"]) - known
        if missing:
            raise ValueError(
                f"'{call['id']}' depends on unknown step(s): {', '.join(sorted(missing))}."
            )
    return calls


def _schedule(calls: list[dict[str, Any]]) -> tuple[list[list[dict[str, Any]]], float]:
    """Topologically group steps and calculate the graph's weighted critical path."""

    by_id = {call["id"]: call for call in calls}
    remaining = set(by_id)
    completed: set[str] = set()
    stages: list[list[dict[str, Any]]] = []
    finish_time: dict[str, float] = {}

    while remaining:
        ready_ids = [
            call["id"]
            for call in calls
            if call["id"] in remaining and set(call["depends_on"]) <= completed
        ]
        if not ready_ids:
            raise ValueError("Dependency cycle detected; the workflow is not a DAG.")
        stage = [by_id[call_id] for call_id in ready_ids]
        stages.append(stage)
        for call in stage:
            start = max((finish_time[item] for item in call["depends_on"]), default=0.0)
            finish_time[call["id"]] = start + call["duration_ms"]
        completed.update(ready_ids)
        remaining.difference_update(ready_ids)

    critical_path_ms = max(finish_time.values(), default=0.0)
    return stages, critical_path_ms


def _metric(label: str, value: str, note: str) -> str:
    return (
        '<div class="metric-card">'
        f'<div class="metric-label">{html.escape(label)}</div>'
        f'<div class="metric-value">{html.escape(value)}</div>'
        f'<div class="metric-note">{html.escape(note)}</div>'
        "</div>"
    )


def _step_card(call: dict[str, Any], index: int | None = None) -> str:
    prefix = f"{index}. " if index is not None else ""
    readonly = "read-only" if call["read_only"] else "side effect"
    return (
        '<div class="step-card">'
        f'<div class="step-tool">{html.escape(prefix + call["tool"])}</div>'
        f'<div class="step-meta">{call["duration_ms"]:,.0f} ms · {readonly}</div>'
        "</div>"
    )


def _render_dashboard(
    trace: dict[str, Any],
    source_calls: list[dict[str, Any]],
    graph_calls: list[dict[str, Any]] | None = None,
) -> str:
    graph_calls = graph_calls or source_calls
    stages, critical_path_ms = _schedule(graph_calls)
    sequential_ms = sum(call["duration_ms"] for call in source_calls)
    theoretical_speedup = sequential_ms / critical_path_ms if critical_path_ms else 1.0
    parallel_steps = sum(max(0, len(stage) - 1) for stage in stages)
    name = html.escape(str(trace.get("name") or "unnamed trace"))

    baseline = "".join(_step_card(call, index + 1) for index, call in enumerate(source_calls))
    compiled_columns = []
    for index, stage in enumerate(stages, start=1):
        stage_duration = max((call["duration_ms"] for call in stage), default=0.0)
        cards = "".join(_step_card(call) for call in stage)
        compiled_columns.append(
            '<div class="stage-column">'
            f'<div class="stage-label">Stage {index} · {stage_duration:,.0f} ms</div>'
            f"{cards}</div>"
        )

    return f"""
    <section class="dashboard" aria-label="Trace compilation result">
      <div class="trace-title"><span class="live-dot"></span>{name}</div>
      <div class="metric-grid">
        {_metric("Source steps", str(len(source_calls)), "read from JSON")}
        {_metric("Compiled stages", str(len(stages)), "topological schedule")}
        {_metric("Sequential latency", f"{sequential_ms / 1000:.2f}s", "sum of declared durations")}
        {_metric("Graph latency", f"{critical_path_ms / 1000:.2f}s", "weighted critical path")}
        {_metric("Parallelizable steps", str(parallel_steps), "same-stage overlap")}
        {
        _metric(
            "Trace-model speedup",
            f"{theoretical_speedup:.2f}×",
            "not a production benchmark",
        )
    }
      </div>
      <div class="comparison-grid">
        <article class="lane baseline-lane">
          <div class="lane-kicker">BEFORE · SOURCE TRACE</div>
          <h3>Sequential replay</h3>
          <div class="baseline-flow">{baseline}</div>
        </article>
        <article class="lane compiled-lane">
          <div class="lane-kicker">AFTER · COMPILED DAG</div>
          <h3>Dependency-aware execution</h3>
          <div class="compiled-flow">{"".join(compiled_columns)}</div>
        </article>
      </div>
      <div class="method-note">
        <strong>How this number is calculated:</strong> sequential latency is the sum of
        <code>duration_ms</code>; graph latency is the weighted critical path allowed by
        <code>depends_on</code>. Tool execution is simulated in this Space. These values
        demonstrate scheduling mechanics and do not prove production gains.
      </div>
    </section>
    """


def _preview_workflow(trace: dict[str, Any], calls: list[dict[str, Any]]) -> dict[str, Any]:
    stages, critical_path_ms = _schedule(calls)
    return {
        "kind": "trace2mcp.workflow.preview",
        "name": trace.get("name", "compiled_workflow"),
        "safety": {
            "allow_side_effects": False,
            "status": "blocked" if any(not call["read_only"] for call in calls) else "read_only",
        },
        "inputs": trace.get("inputs", {}),
        "stages": [
            [
                {
                    "id": call["id"],
                    "tool": call["tool"],
                    "depends_on": call["depends_on"],
                    "duration_ms": call["duration_ms"],
                }
                for call in stage
            ]
            for stage in stages
        ],
        "estimated_critical_path_ms": critical_path_ms,
    }


def _compiled_graph_calls(
    compiled: Any, source_calls: list[dict[str, Any]]
) -> list[dict[str, Any]]:
    """Overlay compiler-derived dependencies onto recorded presentation data."""

    payload = _serializable(compiled)
    if not isinstance(payload, dict) or not isinstance(payload.get("nodes"), list):
        return source_calls
    source_by_id = {call["id"]: call for call in source_calls}
    graph_calls: list[dict[str, Any]] = []
    for raw_node in payload["nodes"]:
        if not isinstance(raw_node, dict) or "id" not in raw_node:
            continue
        node_id = str(raw_node["id"])
        recorded = source_by_id.get(node_id, {})
        graph_calls.append(
            {
                **recorded,
                "id": node_id,
                "tool": str(raw_node.get("tool") or recorded.get("tool") or "tool"),
                "duration_ms": float(recorded.get("duration_ms", 0)),
                "depends_on": [str(item) for item in raw_node.get("depends_on", [])],
                "read_only": raw_node.get("side_effect", "read") == "read",
            }
        )
    return graph_calls or source_calls


def compile_from_text(trace_text: str) -> tuple[str, str, str]:
    """Compile user JSON and return status, visual report, and serialized workflow."""

    try:
        raw = json.loads(trace_text)
        if not isinstance(raw, dict):
            raise ValueError("The top-level JSON value must be an object.")

        parsed: Any = raw
        engine_note = "Preview scheduler active; install the package to compile with Trace2MCP."
        if LOAD_TRACE is not None:
            parsed = _parse_with_package(raw)
        # Keep the source mapping for the visual metrics. Trace v1 deliberately
        # ignores presentation-only fields such as duration_ms/depends_on.
        trace = raw
        calls = _calls_from(trace)

        side_effects = [call["id"] for call in calls if not call["read_only"]]
        if side_effects:
            names = ", ".join(side_effects)
            raise ValueError(
                "Automatic compilation accepts pure and read-only steps only. "
                f"Blocked: {names}. Review them manually or remove them from this demo."
            )

        compiled: Any = _preview_workflow(trace, calls)
        if COMPILE_TRACE is not None:
            compiled = COMPILE_TRACE(parsed, allow_side_effects=False)
            engine_note = "Trace2MCP package compiled this trace with side effects disabled."

        graph_calls = _compiled_graph_calls(compiled, calls)
        output = json.dumps(_serializable(compiled), indent=2, ensure_ascii=False)
        status = (
            "### ✓ Compilation complete\n"
            f"**Engine:** `{ENGINE_MODE}` · {engine_note}  \n"
            "No network calls or API keys were used."
        )
        return status, _render_dashboard(trace, calls, graph_calls), output
    except json.JSONDecodeError as error:
        message = f"Invalid JSON near line {error.lineno}, column {error.colno}: {error.msg}"
    except Exception as error:  # UI boundary: show actionable errors instead of crashing
        message = str(error) or type(error).__name__

    safe_message = html.escape(message)
    return (
        f"### Compilation stopped\n`{safe_message}`",
        (
            '<div class="error-panel"><strong>Could not compile this trace.</strong><br>'
            f"{safe_message}</div>"
        ),
        "{}",
    )


SAMPLE_TEXT = _sample_text()
INITIAL_STATUS, INITIAL_DASHBOARD, INITIAL_OUTPUT = compile_from_text(SAMPLE_TEXT)
CSS = STYLE_PATH.read_text(encoding="utf-8") if STYLE_PATH.exists() else ""


GRADIO_MAJOR = int(gr.__version__.split(".", 1)[0])
BLOCKS_STYLE = {} if GRADIO_MAJOR >= 6 else {"css": CSS, "theme": gr.themes.Soft()}
LAUNCH_STYLE = {"css": CSS, "theme": gr.themes.Soft()} if GRADIO_MAJOR >= 6 else {}


with gr.Blocks(title="Trace2MCP · Teach once. Run forever.", **BLOCKS_STYLE) as demo:
    gr.HTML(
        """
        <header class="hero">
          <div class="eyebrow">OPEN-SOURCE AGENT INFRASTRUCTURE</div>
          <h1>Teach once. <span>Run forever.</span></h1>
          <p>Turn a successful tool-call trace into a deterministic, reusable workflow.</p>
          <div class="hero-badges">
            <span>Offline demo</span><span>No API key</span><span>Read-only by default</span>
          </div>
        </header>
        """
    )
    gr.Markdown(
        "Paste a trace or use the included example. **Compile Trace** validates dependencies, "
        "builds a DAG, and compares sequential replay with dependency-aware scheduling."
    )

    with gr.Row(equal_height=False):
        with gr.Column(scale=5):
            trace_input = gr.Code(
                value=SAMPLE_TEXT,
                language="json",
                label="Agent trace · JSON",
            )
            compile_button = gr.Button("Compile Trace →", variant="primary")
        with gr.Column(scale=4):
            status_output = gr.Markdown(value=INITIAL_STATUS)
            with gr.Accordion("Compiled workflow JSON", open=True):
                workflow_output = gr.Code(
                    value=INITIAL_OUTPUT,
                    language="json",
                    label="Compiler output",
                )

    dashboard_output = gr.HTML(value=INITIAL_DASHBOARD)

    with gr.Accordion("What this demo proves — and what it does not", open=False):
        gr.Markdown(
            """
            **It proves:** a trace can be validated, represented as a dependency graph, and
            scheduled into parallel stages without a model or remote service.

            **It does not prove:** production speedup, semantic equivalence for arbitrary agent
            behavior, or safety of write operations. Durations are supplied by the demo JSON;
            real deployments must benchmark real tools and review generated workflows.
            """
        )

    gr.HTML('<footer class="footer">Trace2MCP prototype · War Enterprise · MIT</footer>')

    compile_button.click(
        fn=compile_from_text,
        inputs=trace_input,
        outputs=[status_output, dashboard_output, workflow_output],
    )


if __name__ == "__main__":
    demo.launch(**LAUNCH_STYLE)