"""EB1-GNN Conductor: Intelligent inference with LLM-guided orchestration.

This module implements a two-tier inference architecture:

Tier 1 - Fast Path (GNN Direct):
    For queries where the GNN is confident, route directly to the predicted
    best model. No LLM conductor overhead.

Tier 2 - Conductor Path (LLM + GNN Tool):
    For hard/uncertain queries, an LLM conductor orchestrates the response.
    The conductor can:
    - Decompose complex tasks into subtasks
    - Invoke the GNN router as a tool to select models for each subtask
    - Synthesize final answers from subtask outputs

The routing decision uses signals extracted directly from GNN logits:
- confidence: max(softmax(logits)) - how certain is the top prediction
- entropy: normalized Shannon entropy - how spread is the distribution
- difficulty: 1 - confidence - intrinsic query difficulty estimate

Threshold Design:
    - TRIVIAL: confidence > 0.7 AND entropy < 0.3 -> cheap model (any will work)
    - EASY: confidence > 0.5 -> GNN direct routing
    - HARD: confidence <= 0.5 OR entropy > 0.6 -> conductor path

Example:
    from ember.takumi.eb1_gnn_conductor import EB1GNNInference

    # Initialize with default thresholds
    inference = EB1GNNInference.from_registry("eb1-gnn-preview-12-19")

    # Route and execute
    result = inference("What is the capital of France?")  # Fast path
    result = inference("Derive the Navier-Stokes equations")  # Conductor path

    # Access routing metadata
    print(result.routing_path)  # "fast" or "conductor"
    print(result.routing_signals)  # RoutingSignals dataclass
"""

from __future__ import annotations

import json
import logging
import re
import threading
import time
from collections import deque
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import Any, Protocol, runtime_checkable

from pydantic import BaseModel, ConfigDict, Field

from ember.api import models
from ember.api.models import ModelBinding
from ember.models.schemas import ChatResponse, UsageStats
from ember.operators.base import Operator
from .model_pool import CLAUDE_OPUS, GNN_LABEL_TO_MODEL_ID

_CONDUCTOR_LOG = logging.getLogger("ember.conductor")


# =============================================================================
# GNN Router Protocol (structural typing for router abstraction)
# =============================================================================


@runtime_checkable
class GNNRouteResult(Protocol):
    """Structural protocol for GNN routing results.

    Any object with these attributes satisfies this protocol, including
    MockRoutingResult in tests and EB1Router results in production.
    """

    best_model: str
    best_model_idx: int
    confidence: float
    entropy: float
    probabilities: dict[str, float]
    scores: dict[str, float]
    total_latency_ms: float


@runtime_checkable
class GNNRouterProtocol(Protocol):
    """Structural protocol for GNN router instances.

    Follows the EmbeddingProvider pattern from routing.py: runtime-checkable
    so callers can validate at construction time if needed.

    Implementors:
        - EB1Router (production, from embergraphrouter.inference)
        - TieredRouter (production, from embergraphrouter.inference)
        - MockGNNRouter (testing)
    """

    def route(self, query: str) -> GNNRouteResult:
        """Route a query and return model selection with confidence signals.

        Args:
            query: The text to route.

        Returns:
            Result with best_model, confidence, entropy, probabilities,
            scores, and total_latency_ms.
        """
        ...

# XML extraction for conductor directives, matching the proven <response> tag
# pattern from eb1.py. Falls back to bare JSON extraction on miss.
_DIRECTIVE_TAG_RE = re.compile(r"<directive>(.*?)</directive>", re.DOTALL)


# =============================================================================
# Routing Signals and Thresholds
# =============================================================================


class RoutingTier(str, Enum):
    """Tier classification for routing decisions."""

    TRIVIAL = "trivial"  # All models will succeed - use cheapest
    EASY = "easy"  # GNN confident - use direct routing
    HARD = "hard"  # GNN uncertain - use conductor


@dataclass(frozen=True)
class RoutingThresholds:
    """Threshold configuration for tier classification.

    These thresholds are derived from empirical analysis of the GNN router's
    calibration on the training distribution. Key observations:

    1. When confidence > 0.7, the GNN's top-1 prediction is almost always
       correct (precision > 90%). These are "easy" queries.

    2. When entropy > 0.6, the GNN can't distinguish between models well.
       This indicates either a trivial query (all models succeed) or a
       genuinely hard query (no clear winner).

    3. Mean logit (average across all model predictions) correlates with
       difficulty: high mean = trivial, low mean = hard.

    The conductor escalation threshold balances cost vs quality:
    - Lower threshold = more conductor usage = higher quality, higher cost
    - Higher threshold = less conductor usage = lower quality, lower cost
    """

    # Trivial detection (all models likely to succeed)
    trivial_entropy_max: float = 0.3  # Low entropy = uniform success
    trivial_confidence_min: float = 0.7  # But still confident

    # Easy detection (GNN can route directly)
    easy_confidence_min: float = 0.5  # Trust GNN top-1

    # Conductor escalation (GNN uncertain)
    conductor_entropy_min: float = 0.6  # High entropy = uncertain
    conductor_confidence_max: float = 0.5  # Low confidence = escalate

    # Mean logit thresholds (from difficulty detection research)
    trivial_mean_logit_min: float = 0.5  # High mean = likely trivial
    hard_mean_logit_max: float = -0.3  # Low mean = likely hard


# Module-level singleton avoids allocating a new RoutingThresholds on every
# RoutingSignals.tier access (frozen dataclass, safe to share).
_DEFAULT_THRESHOLDS = RoutingThresholds()


@dataclass
class RoutingSignals:
    """Signals extracted from GNN predictions for routing decisions.

    These signals are computed directly from the GNN's output logits without
    requiring any additional trainable parameters. The key insight is that
    a well-calibrated router's uncertainty (entropy) correlates with query
    difficulty.
    """

    confidence: float  # max(softmax(logits))
    entropy: float  # Normalized Shannon entropy [0, 1]
    difficulty: float  # 1 - confidence
    mean_logit: float  # Average logit (difficulty proxy)
    logit_gap: float  # top1 - top2 logit (discrimination)
    chosen_model: str  # GNN's top-1 prediction
    chosen_model_idx: int
    probabilities: dict[str, float]  # All model probabilities

    @property
    def tier(self) -> RoutingTier:
        """Classify query into routing tier using default thresholds."""
        return self.classify_with_thresholds(_DEFAULT_THRESHOLDS)

    def classify_with_thresholds(self, thresholds: RoutingThresholds) -> RoutingTier:
        """Classify with custom thresholds."""
        if (
            self.confidence >= thresholds.trivial_confidence_min
            and self.entropy <= thresholds.trivial_entropy_max
        ):
            return RoutingTier.TRIVIAL

        if self.confidence >= thresholds.easy_confidence_min:
            return RoutingTier.EASY

        return RoutingTier.HARD


# =============================================================================
# GNN Router Tool (for Conductor LLM)
# =============================================================================


class GNNRoutingTool:
    """Wraps GNN router as a tool callable by the conductor LLM.

    The conductor can invoke this tool to get model recommendations for
    specific queries or subtasks. This enables the conductor to:
    1. Query the GNN for model selection on decomposed subtasks
    2. Get confidence/uncertainty estimates to inform its decisions
    3. Access the full probability distribution over models

    Example tool invocation from conductor:
        {
            "tool": "gnn_router",
            "query": "Compute the integral of sin(x)dx",
            "return_top_k": 3
        }

    Returns:
        {
            "recommended_model": "Claude Opus 4.5",
            "confidence": 0.72,
            "uncertainty": "low",
            "top_models": [
                {"model": "Claude Opus 4.5", "probability": 0.72},
                {"model": "GPT-5", "probability": 0.18},
                {"model": "Gemini 3 Pro", "probability": 0.06}
            ]
        }
    """

    def __init__(self, gnn_router: GNNRouterProtocol):
        """Initialize with GNN router instance.

        Args:
            gnn_router: Any object implementing the GNNRouterProtocol.
        """
        self._router = gnn_router

    def __call__(
        self,
        query: str,
        return_top_k: int = 3,
    ) -> dict[str, Any]:
        """Route a query and return structured recommendation.

        Args:
            query: The query or subtask to route.
            return_top_k: Number of top models to return.

        Returns:
            Dictionary with routing recommendation and metadata.
        """
        result = self._router.route(query)

        # Sort models by probability
        sorted_models = sorted(
            result.probabilities.items(),
            key=lambda x: x[1],
            reverse=True,
        )

        # Classify uncertainty level
        if result.confidence > 0.7:
            uncertainty_level = "low"
        elif result.confidence > 0.5:
            uncertainty_level = "medium"
        else:
            uncertainty_level = "high"

        return {
            "recommended_model": result.best_model,
            "confidence": round(result.confidence, 3),
            "uncertainty": uncertainty_level,
            "entropy": round(result.entropy, 3),
            "top_models": [
                {"model": name, "probability": round(prob, 3)}
                for name, prob in sorted_models[:return_top_k]
            ],
            "latency_ms": round(result.total_latency_ms, 2),
        }

    def get_routing_signals(self, query: str) -> RoutingSignals:
        """Get full routing signals for a query.

        Args:
            query: The query to analyze.

        Returns:
            RoutingSignals with all computed metrics.
        """
        result = self._router.route(query)

        # Compute mean logit (need access to scores/logits)
        scores = list(result.scores.values())
        mean_logit = sum(scores) / len(scores)

        # Compute logit gap
        sorted_scores = sorted(scores, reverse=True)
        logit_gap = sorted_scores[0] - sorted_scores[1] if len(sorted_scores) > 1 else 0.0

        return RoutingSignals(
            confidence=result.confidence,
            entropy=result.entropy,
            difficulty=1.0 - result.confidence,
            mean_logit=mean_logit,
            logit_gap=logit_gap,
            chosen_model=result.best_model,
            chosen_model_idx=result.best_model_idx,
            probabilities=result.probabilities,
        )


# =============================================================================
# Metrics Tracking
# =============================================================================


@dataclass
class GNNCallRecord:
    """Record of a single GNN router invocation.

    Tracks what query was routed, what the GNN recommended, and how
    confident it was. Essential for understanding GNN utilization patterns.
    """

    query: str
    query_truncated: str  # First 100 chars for logging
    recommended_model: str
    confidence: float
    entropy: float
    tier: RoutingTier
    probabilities: dict[str, float]
    timestamp_ms: float
    latency_ms: float
    caller: str  # "initial", "subtask", "explicit_route"

    @classmethod
    def from_signals(
        cls,
        query: str,
        signals: "RoutingSignals",
        latency_ms: float,
        caller: str,
    ) -> "GNNCallRecord":
        """Create a record from routing signals."""
        return cls(
            query=query,
            query_truncated=query[:100] + "..." if len(query) > 100 else query,
            recommended_model=signals.chosen_model,
            confidence=signals.confidence,
            entropy=signals.entropy,
            tier=signals.tier,
            probabilities=signals.probabilities.copy(),
            timestamp_ms=time.time() * 1000,
            latency_ms=latency_ms,
            caller=caller,
        )


@dataclass
class ConductorMetrics:
    """Comprehensive metrics for conductor decision-making.

    Aggregates all metrics from a session of inference calls, enabling
    analysis of:
    - GNN utilization: How often is the GNN called? For what?
    - Decision distribution: How often does conductor choose each action?
    - Cost analysis: What's the estimated cost breakdown?
    - Latency breakdown: Where is time spent?
    """

    # GNN call tracking
    gnn_calls: list[GNNCallRecord] = field(default_factory=list)

    # Path distribution
    queries_trivial: int = 0
    queries_easy: int = 0
    queries_hard: int = 0

    # Conductor action distribution (for hard queries)
    action_counts: dict[str, int] = field(
        default_factory=lambda: {
            "direct": 0,
            "decompose": 0,
            "synthesize": 0,
            "delegate": 0,
        }
    )

    # Subtask statistics
    total_subtasks: int = 0
    subtasks_with_gnn_routing: int = 0
    subtasks_with_override: int = 0

    # Model usage
    model_usage_counts: dict[str, int] = field(default_factory=dict)

    # Latency tracking
    total_latency_ms: float = 0.0
    gnn_latency_ms: float = 0.0
    conductor_latency_ms: float = 0.0
    execution_latency_ms: float = 0.0

    # Cost estimation (relative units)
    estimated_cost: float = 0.0

    def record_gnn_call(self, record: GNNCallRecord) -> None:
        """Record a GNN invocation."""
        self.gnn_calls.append(record)
        self.gnn_latency_ms += record.latency_ms

    def record_path(self, tier: RoutingTier) -> None:
        """Record which path a query took."""
        if tier == RoutingTier.TRIVIAL:
            self.queries_trivial += 1
        elif tier == RoutingTier.EASY:
            self.queries_easy += 1
        else:
            self.queries_hard += 1

    def record_action(self, action: str) -> None:
        """Record a conductor action."""
        if action in self.action_counts:
            self.action_counts[action] += 1

    def record_model_usage(self, model: str) -> None:
        """Record model usage."""
        self.model_usage_counts[model] = self.model_usage_counts.get(model, 0) + 1

    def record_subtask(self, used_gnn: bool, had_override: bool) -> None:
        """Record subtask execution."""
        self.total_subtasks += 1
        if used_gnn:
            self.subtasks_with_gnn_routing += 1
        if had_override:
            self.subtasks_with_override += 1

    @property
    def total_queries(self) -> int:
        return self.queries_trivial + self.queries_easy + self.queries_hard

    @property
    def gnn_call_count(self) -> int:
        return len(self.gnn_calls)

    @property
    def avg_gnn_confidence(self) -> float:
        if not self.gnn_calls:
            return 0.0
        return sum(c.confidence for c in self.gnn_calls) / len(self.gnn_calls)

    @property
    def avg_gnn_entropy(self) -> float:
        if not self.gnn_calls:
            return 0.0
        return sum(c.entropy for c in self.gnn_calls) / len(self.gnn_calls)

    def path_distribution(self) -> dict[str, float]:
        """Get percentage distribution across paths."""
        total = self.total_queries
        if total == 0:
            return {"trivial": 0.0, "easy": 0.0, "hard": 0.0}
        return {
            "trivial": self.queries_trivial / total,
            "easy": self.queries_easy / total,
            "hard": self.queries_hard / total,
        }

    def gnn_calls_by_caller(self) -> dict[str, int]:
        """Count GNN calls by caller type."""
        counts: dict[str, int] = {}
        for call in self.gnn_calls:
            counts[call.caller] = counts.get(call.caller, 0) + 1
        return counts

    def to_summary_dict(self) -> dict[str, Any]:
        """Generate summary dictionary for logging/analysis."""
        return {
            "total_queries": self.total_queries,
            "path_distribution": self.path_distribution(),
            "gnn_calls": {
                "total": self.gnn_call_count,
                "by_caller": self.gnn_calls_by_caller(),
                "avg_confidence": round(self.avg_gnn_confidence, 3),
                "avg_entropy": round(self.avg_gnn_entropy, 3),
            },
            "conductor_actions": self.action_counts,
            "subtasks": {
                "total": self.total_subtasks,
                "with_gnn_routing": self.subtasks_with_gnn_routing,
                "with_override": self.subtasks_with_override,
            },
            "model_usage": self.model_usage_counts,
            "latency_breakdown_ms": {
                "total": round(self.total_latency_ms, 2),
                "gnn": round(self.gnn_latency_ms, 2),
                "conductor": round(self.conductor_latency_ms, 2),
                "execution": round(self.execution_latency_ms, 2),
            },
            "estimated_cost": round(self.estimated_cost, 2),
        }

    def print_summary(self) -> None:
        """Print a human-readable summary."""
        summary = self.to_summary_dict()

        print("\n" + "=" * 70)
        print("CONDUCTOR METRICS SUMMARY")
        print("=" * 70)

        print(f"\nTotal Queries: {summary['total_queries']}")
        print("Path Distribution:")
        for path, pct in summary["path_distribution"].items():
            print(f"  {path}: {pct:.1%}")

        print(f"\nGNN Calls: {summary['gnn_calls']['total']}")
        print("  By caller:")
        for caller, count in summary["gnn_calls"]["by_caller"].items():
            print(f"    {caller}: {count}")
        print(f"  Avg confidence: {summary['gnn_calls']['avg_confidence']:.3f}")
        print(f"  Avg entropy: {summary['gnn_calls']['avg_entropy']:.3f}")

        if self.queries_hard > 0:
            print("\nConductor Actions:")
            for action, count in summary["conductor_actions"].items():
                if count > 0:
                    print(f"  {action}: {count}")

            if summary["subtasks"]["total"] > 0:
                print(f"\nSubtasks: {summary['subtasks']['total']}")
                print(f"  With GNN routing: {summary['subtasks']['with_gnn_routing']}")
                print(f"  With model override: {summary['subtasks']['with_override']}")

        if summary["model_usage"]:
            print("\nModel Usage:")
            for model, count in sorted(
                summary["model_usage"].items(), key=lambda x: x[1], reverse=True
            ):
                print(f"  {model}: {count}")

        print("\nLatency Breakdown:")
        for component, ms in summary["latency_breakdown_ms"].items():
            print(f"  {component}: {ms:.1f}ms")

        print("=" * 70)


# =============================================================================
# Conductor Schemas
# =============================================================================


class ConductorAction(str, Enum):
    """Actions the conductor can take.

    DIRECT: Call a specific model with the full query.
    DECOMPOSE: Break query into parallel subtasks for different models.
    SYNTHESIZE: Combine subtask results into a final answer.
    DELEGATE: Delegate to the full EB1 multi-model ensemble.
    """

    DIRECT = "direct"
    DECOMPOSE = "decompose"
    SYNTHESIZE = "synthesize"
    DELEGATE = "delegate"


class SubtaskSpec(BaseModel):
    """Specification for a subtask created by decomposition."""

    model_config = ConfigDict(validate_by_name=True)

    id: str = Field(description="Unique subtask identifier")
    description: str = Field(description="What this subtask accomplishes")
    query: str = Field(description="The specific query for this subtask")
    depends_on: list[str] = Field(
        default_factory=list,
        description="IDs of subtasks this depends on",
    )
    preferred_model: str | None = Field(
        default=None,
        description="Optional model override (else use GNN)",
    )


class ConductorDirective(BaseModel):
    """Schema for conductor's action directives."""

    model_config = ConfigDict(validate_by_name=True)

    action: ConductorAction
    reasoning: str = Field(description="Why this action was chosen")

    # For DIRECT action
    target_model: str | None = None

    # For DECOMPOSE action
    subtasks: list[SubtaskSpec] | None = None

    # For SYNTHESIZE action
    synthesis_prompt: str | None = None

    # For DELEGATE action
    delegate_config: str | None = None  # "eb1" or "eb1-pro"

    # Control flow
    is_final: bool = Field(default=False)
    final_answer: str | None = None


@dataclass
class SubtaskResult:
    """Result from executing a subtask."""

    subtask_id: str
    model_used: str
    response: str
    latency_ms: float
    routing_signals: RoutingSignals | None = None
    usage: UsageStats | None = None
    error: Exception | None = None  # Track errors without failing execution


# =============================================================================
# Budget and Parallelism Tracking
# =============================================================================


@dataclass(frozen=True)
class ConductorBudget:
    """Budget limits for conductor orchestration.

    Enforces resource constraints across multiple dimensions to prevent
    runaway orchestration costs. All limits are optional (None = unlimited).

    Attributes:
        max_tokens: Total token budget (prompt + completion) for entire forward() call.
        max_conductor_calls: Maximum conductor LLM invocations per forward().
        max_subtasks: Maximum total subtasks that can be executed.
        max_latency_ms: Time limit for entire forward() call (wall-clock time).
        max_subtask_latency_ms: Per-subtask execution time limit.
    """

    max_tokens: int | None = None
    max_conductor_calls: int | None = None
    max_subtasks: int | None = None
    max_latency_ms: float | None = None
    max_subtask_latency_ms: float | None = None

    def __post_init__(self) -> None:
        """Validate budget limits are positive."""
        for field_name in ["max_tokens", "max_conductor_calls", "max_subtasks"]:
            value = getattr(self, field_name)
            if value is not None and value <= 0:
                raise ValueError(f"{field_name} must be positive, got {value}")

        for field_name in ["max_latency_ms", "max_subtask_latency_ms"]:
            value = getattr(self, field_name)
            if value is not None and value <= 0.0:
                raise ValueError(f"{field_name} must be positive, got {value}")


class BudgetExhaustedError(Exception):
    """Raised when conductor budget is exhausted."""

    def __init__(self, dimension: str, limit: int | float, consumed: int | float):
        self.dimension = dimension
        self.limit = limit
        self.consumed = consumed
        super().__init__(
            f"Budget exhausted: {dimension} limit={limit}, consumed={consumed}"
        )


@dataclass
class BudgetTracker:
    """Thread-safe tracker for budget consumption during conductor execution."""

    budget: ConductorBudget
    consumed_tokens: int = 0
    conductor_calls: int = 0
    subtasks_executed: int = 0
    start_time_ms: float = field(default_factory=lambda: time.perf_counter() * 1000)
    _lock: threading.Lock = field(default_factory=threading.Lock, repr=False)

    def check_and_increment_conductor_call(self) -> None:
        """Atomically check and increment conductor call counter."""
        with self._lock:
            if self.budget.max_conductor_calls is not None:
                if self.conductor_calls >= self.budget.max_conductor_calls:
                    raise BudgetExhaustedError(
                        dimension="conductor_calls",
                        limit=self.budget.max_conductor_calls,
                        consumed=self.conductor_calls,
                    )
            self.conductor_calls += 1

    def check_and_increment_subtask(self) -> None:
        """Atomically check and increment subtask counter."""
        with self._lock:
            if self.budget.max_subtasks is not None:
                if self.subtasks_executed >= self.budget.max_subtasks:
                    raise BudgetExhaustedError(
                        dimension="subtasks",
                        limit=self.budget.max_subtasks,
                        consumed=self.subtasks_executed,
                    )
            self.subtasks_executed += 1

    def check_tokens(self, additional_tokens: int) -> None:
        """Check if token budget allows additional consumption."""
        with self._lock:
            if self.budget.max_tokens is not None:
                if self.consumed_tokens + additional_tokens > self.budget.max_tokens:
                    raise BudgetExhaustedError(
                        dimension="tokens",
                        limit=self.budget.max_tokens,
                        consumed=self.consumed_tokens + additional_tokens,
                    )

    def consume_tokens(self, tokens: int) -> None:
        """Record token consumption."""
        with self._lock:
            self.consumed_tokens += tokens

    def check_latency(self) -> None:
        """Check if latency budget exceeded."""
        if self.budget.max_latency_ms is not None:
            elapsed_ms = time.perf_counter() * 1000 - self.start_time_ms
            if elapsed_ms >= self.budget.max_latency_ms:
                raise BudgetExhaustedError(
                    dimension="latency_ms",
                    limit=self.budget.max_latency_ms,
                    consumed=elapsed_ms,
                )

    def to_dict(self) -> dict[str, Any]:
        """Export budget consumption as dictionary."""
        elapsed_ms = time.perf_counter() * 1000 - self.start_time_ms
        return {
            "consumed_tokens": self.consumed_tokens,
            "conductor_calls": self.conductor_calls,
            "subtasks_executed": self.subtasks_executed,
            "elapsed_ms": elapsed_ms,
            "limits": {
                "max_tokens": self.budget.max_tokens,
                "max_conductor_calls": self.budget.max_conductor_calls,
                "max_subtasks": self.budget.max_subtasks,
                "max_latency_ms": self.budget.max_latency_ms,
            },
        }


@dataclass
class ParallelismMetrics:
    """Metrics for parallel subtask execution effectiveness."""

    # Per-subtask latency tracking
    subtask_latencies_ms: dict[str, float] = field(default_factory=dict)

    # Wave-based execution tracking
    execution_waves: list[list[str]] = field(default_factory=list)
    wave_latencies_ms: list[float] = field(default_factory=list)

    # Parallelism efficiency
    total_parallel_latency_ms: float = 0.0  # Sum of wave latencies
    total_sequential_latency_ms: float = 0.0  # Sum of subtask latencies
    actual_speedup: float = 1.0  # sequential / parallel
    theoretical_max_speedup: float = 1.0  # Based on critical path

    # Critical path analysis
    critical_path: list[str] = field(default_factory=list)
    critical_path_latency_ms: float = 0.0

    # Error tracking
    failed_subtasks: list[str] = field(default_factory=list)
    partial_synthesis: bool = False

    def compute_speedup(self) -> None:
        """Compute actual speedup from sequential vs parallel execution."""
        if self.total_parallel_latency_ms > 0:
            self.actual_speedup = (
                self.total_sequential_latency_ms / self.total_parallel_latency_ms
            )
        if self.critical_path_latency_ms > 0:
            self.theoretical_max_speedup = (
                self.total_sequential_latency_ms / self.critical_path_latency_ms
            )


@dataclass
class ConductorTrace:
    """Trace of conductor's decision process."""

    directives: list[ConductorDirective] = field(default_factory=list)
    subtask_results: list[SubtaskResult] = field(default_factory=list)
    total_latency_ms: float = 0.0
    conductor_calls: int = 0

    # New fields for budget and parallelism tracking
    budget_consumption: dict[str, Any] | None = None
    parallelism_metrics: ParallelismMetrics | None = None
    budget_exhausted: bool = False
    budget_exhausted_dimension: str | None = None

    def aggregate_usage(self) -> UsageStats:
        """Sum token usage across all subtask results.

        Returns:
            Aggregated UsageStats with totals from all subtasks that
            reported usage. Subtasks with ``usage=None`` are skipped.
        """
        prompt = 0
        completion = 0
        for result in self.subtask_results:
            if result.usage is not None:
                prompt += result.usage.prompt_tokens
                completion += result.usage.completion_tokens
        return UsageStats(
            prompt_tokens=prompt,
            completion_tokens=completion,
        )


# =============================================================================
# EB1 Conductor Operator
# =============================================================================


# Default model pool — extends canonical GNN label mapping with
# conductor-specific models (RND1).  Core 8-LLM set lives in model_pool.py.
_DEFAULT_MODEL_POOL: dict[str, str] = {
    **GNN_LABEL_TO_MODEL_ID,
    "RND1 Base": "rnd1/rnd1-base-0910",
}


class EB1Conductor(Operator):
    """LLM conductor that orchestrates complex queries using GNN as a tool.

    The conductor receives hard queries that the GNN couldn't confidently
    route. It can:

    1. Decompose the query into subtasks
    2. Use the GNN router tool to select models for each subtask
    3. Execute subtasks (potentially in parallel for independent ones)
    4. Synthesize results into a final answer

    For very complex queries, it can delegate to the full EB1 ensemble.

    Attributes:
        conductor_model: The LLM that makes orchestration decisions.
        gnn_tool: GNN router wrapped as a callable tool.
        model_pool: Available models for execution.
    """

    conductor_model: ModelBinding
    gnn_tool: GNNRoutingTool
    model_pool: dict[str, ModelBinding]
    default_model: str
    max_iterations: int

    def __init__(
        self,
        gnn_router: GNNRouterProtocol,
        conductor_model: str = CLAUDE_OPUS,
        model_pool: dict[str, str] | None = None,
        max_iterations: int = 5,
        budget: ConductorBudget | None = None,
        default_model: str = "Claude Opus 4.5",
    ):
        """Initialize the conductor.

        Args:
            gnn_router: Any object implementing GNNRouterProtocol.
            conductor_model: Model ID for the orchestration LLM.
            model_pool: Mapping of model names to IDs. If None, uses defaults.
            max_iterations: Maximum conductor iterations before forcing completion.
            budget: Budget constraints for conductor execution. If None, no limits.
            default_model: Display name of the default/strongest model in pool.
                Must be a key in model_pool.
        """
        object.__setattr__(self, "conductor_model", models.instance(conductor_model))
        object.__setattr__(self, "gnn_tool", GNNRoutingTool(gnn_router))
        object.__setattr__(self, "max_iterations", max_iterations)
        object.__setattr__(self, "budget", budget or ConductorBudget())

        pool_config = model_pool or _DEFAULT_MODEL_POOL
        object.__setattr__(
            self,
            "model_pool",
            {name: models.instance(model_id) for name, model_id in pool_config.items()},
        )

        # Validate default_model exists in pool; fall back to first key if not.
        if default_model in pool_config:
            object.__setattr__(self, "default_model", default_model)
        else:
            first_key = next(iter(pool_config))
            object.__setattr__(self, "default_model", first_key)

    def forward(
        self,
        query: str,
        *,
        system: str | None = None,
        routing_signals: RoutingSignals | None = None,
    ) -> tuple[ChatResponse, ConductorTrace]:
        """Orchestrate a complex query.

        Args:
            query: The user's query.
            system: Optional system prompt.
            routing_signals: Pre-computed routing signals (for context).

        Returns:
            Tuple of (final response, execution trace).
        """
        trace = ConductorTrace()
        subtask_outputs: dict[str, SubtaskResult] = {}
        budget_tracker = BudgetTracker(self.budget)

        # Build initial context for conductor
        context = self._build_conductor_context(query, system, routing_signals)
        conductor_system = self._conductor_system_prompt(self.budget)

        iteration = 0
        final_response: ChatResponse | None = None

        try:
            while iteration < self.max_iterations:
                iteration += 1
                trace.conductor_calls += 1

                # Check budget before conductor call
                budget_tracker.check_and_increment_conductor_call()

                # Get conductor's directive
                start = time.perf_counter()
                directive = self._invoke_conductor(context, conductor_system)
                trace.directives.append(directive)
                trace.total_latency_ms += (time.perf_counter() - start) * 1000

                # Execute the action
                if directive.action == ConductorAction.DIRECT:
                    final_response = self._execute_direct(query, system, directive.target_model, trace)
                    break

                elif directive.action == ConductorAction.DECOMPOSE:
                    if directive.subtasks:
                        # Replace sequential loop with parallel execution
                        new_results = self._execute_subtasks_parallel(
                            directive.subtasks, subtask_outputs, trace, budget_tracker
                        )
                        subtask_outputs.update(new_results)
                        # Update context with results
                        context = self._update_context_with_results(context, directive, subtask_outputs)

                elif directive.action == ConductorAction.SYNTHESIZE:
                    final_response = self._synthesize_results(
                        query, system, subtask_outputs, directive.synthesis_prompt, trace
                    )
                    break

                elif directive.action == ConductorAction.DELEGATE:
                    # Delegate to EB1 ensemble (import here to avoid circular)
                    final_response = self._delegate_to_eb1(
                        query, system, directive.delegate_config, trace
                    )
                    break

                # Check for final answer in directive
                if directive.is_final and directive.final_answer:
                    final_response = ChatResponse(data=directive.final_answer)
                    break

        except BudgetExhaustedError as e:
            _CONDUCTOR_LOG.warning(
                "Conductor budget exhausted (dimension=%s, limit=%s, "
                "consumed=%s) — synthesizing from %d partial subtask outputs",
                e.dimension,
                e.limit,
                e.consumed,
                len(subtask_outputs),
            )
            trace.budget_exhausted = True
            trace.budget_exhausted_dimension = e.dimension
            final_response = self._synthesize_partial(query, system, subtask_outputs, trace)

        finally:
            # Always record budget consumption
            trace.budget_consumption = budget_tracker.to_dict()

        # Fallback if we hit max iterations
        if final_response is None:
            final_response = self._fallback_response(query, system, trace)

        return final_response, trace

    def _build_conductor_context(
        self,
        query: str,
        system: str | None,
        routing_signals: RoutingSignals | None,
    ) -> str:
        """Build initial context for the conductor."""
        context_parts = [
            "<<QUERY>>",
            query,
            "<<END_QUERY>>",
        ]

        if system:
            context_parts.extend(["<<SYSTEM>>", system, "<<END_SYSTEM>>"])

        if routing_signals:
            context_parts.extend(
                [
                    "<<GNN_SIGNALS>>",
                    json.dumps(
                        {
                            "recommended_model": routing_signals.chosen_model,
                            "confidence": round(routing_signals.confidence, 3),
                            "entropy": round(routing_signals.entropy, 3),
                            "difficulty_estimate": round(routing_signals.difficulty, 3),
                            "routing_tier": routing_signals.tier.value,
                            "top_probabilities": {
                                k: round(v, 3)
                                for k, v in sorted(
                                    routing_signals.probabilities.items(),
                                    key=lambda x: x[1],
                                    reverse=True,
                                )[:3]
                            },
                        }
                    ),
                    "<<END_GNN_SIGNALS>>",
                ]
            )

        context_parts.extend(
            [
                "<<AVAILABLE_MODELS>>",
                self._render_model_capabilities(),
                "<<END_AVAILABLE_MODELS>>",
            ]
        )

        return "\n".join(context_parts)

    def _render_model_capabilities(self) -> str:
        """Render model pool with capability annotations for the conductor.

        Returns a compact summary so the conductor knows which models are
        suitable for code, tools/JSON, free-form generation, etc.
        """
        # Static capability map for the GNN 8-LLM training set.
        # Keys are model display names (must match model_pool keys).
        _CAPS: dict[str, str] = {
            "Gemini 2.5 Flash Lite": "fast, cheap, general",
            "Claude Haiku 4.5": "fast, tools, JSON, general",
            "GPT-5 Nano": "fast, cheap, general",
            "GPT-5.2": "strong general, tools, JSON, code, reasoning",
            "GPT-5.2 Codex": "code-specialized, tools, JSON, reasoning",
            "Claude Opus 4.6": (
                "strongest general, tools, JSON, code, reasoning"
            ),
            "Gemini 3 Flash Preview": "fast, tools, JSON, general",
            "Gemini 3.1 Pro Preview": (
                "strong general, tools, JSON, code, reasoning"
            ),
            "RND1 Base": (
                "free-form generation ONLY"
                " (no tools, no JSON, no system prompt)"
            ),
            # Legacy aliases (v15/v14 checkpoints)
            "GPT-5.1": "strong general, tools, JSON, code",
            "GPT-5.1 Codex": "code-specialized, tools, JSON",
            "Claude Opus 4.5": (
                "strongest general, tools, JSON, code, reasoning"
            ),
            "Gemini 3 Pro Preview": (
                "strong general, tools, JSON, code"
            ),
        }
        lines: list[str] = []
        for name in self.model_pool:
            caps = _CAPS.get(name, "general")
            lines.append(f"- {name}: {caps}")
        return "\n".join(lines)

    def _conductor_system_prompt(self, budget: ConductorBudget) -> str:
        """Build the conductor system prompt with budget and capability context.

        Args:
            budget: Current budget constraints (injected so the conductor
                can reason about remaining resources).

        Returns:
            Complete system prompt string.
        """
        # Budget line: only shown when finite limits exist.
        budget_parts: list[str] = []
        if budget.max_subtasks is not None:
            budget_parts.append(f"max_subtasks={budget.max_subtasks}")
        if budget.max_conductor_calls is not None:
            budget_parts.append(f"max_conductor_calls={budget.max_conductor_calls}")
        budget_line = ""
        if budget_parts:
            budget_line = (
                f"\n\nBudget: {', '.join(budget_parts)}. "
                "Prefer DIRECT when uncertain about decomposition value."
            )

        return f"""\
You are the EB1-GNN Conductor, an orchestration agent that handles \
queries where the GNN router is uncertain about the best model.

You receive GNN routing signals (confidence, entropy, model probabilities) \
in the context. Use these signals to inform your decisions.

ACTIONS (exactly 4):
  DIRECT     - Route the full query to a single model.
  DECOMPOSE  - Split the query into independent or dependent subtasks.
  SYNTHESIZE - Combine completed subtask results into a final answer.
  DELEGATE   - Delegate to the full EB1 multi-model ensemble (expensive).

WHEN NOT TO DECOMPOSE:
  - The query is a single coherent task with no separable parts.
  - Decomposition would add latency without quality gain.
  - You are uncertain how to decompose effectively.
  In these cases, prefer DIRECT to the strongest relevant model.

RESPONSE FORMAT:
Wrap your output in <directive> tags containing valid JSON:

<directive>
{{"action": "...", "reasoning": "...", ...}}
</directive>

Text outside the tags is private reasoning (not shown to users).

FEW-SHOT EXAMPLES:

Example 1 (DIRECT - simple factual query):
<directive>
{{"action": "direct", "reasoning": "Single factual question, GNN recommends Claude Opus 4.5 with 0.72 confidence.", "target_model": "Claude Opus 4.5"}}
</directive>

Example 2 (DECOMPOSE - multi-part analysis):
<directive>
{{"action": "decompose", "reasoning": "Query has two independent parts: data analysis and code generation.", "subtasks": [{{"id": "analyze", "description": "Analyze the dataset", "query": "Analyze the trends in the provided sales data", "depends_on": [], "preferred_model": null}}, {{"id": "code", "description": "Generate visualization code", "query": "Write Python matplotlib code to plot sales trends", "depends_on": ["analyze"], "preferred_model": "GPT-5.1 Codex"}}]}}
</directive>

Example 3 (SYNTHESIZE - after subtask results are available):
<directive>
{{"action": "synthesize", "reasoning": "Both subtasks completed. Combining analysis with visualization code.", "synthesis_prompt": "Combine the data analysis findings with the generated visualization code into a coherent report."}}
</directive>

Example 4 (DELEGATE - requires multi-model deliberation):
<directive>
{{"action": "delegate", "reasoning": "Highly subjective creative task benefits from multi-model ensemble.", "delegate_config": "eb1-pro"}}
</directive>{budget_line}"""

    def _invoke_conductor(self, context: str, system: str) -> ConductorDirective:
        """Invoke the conductor LLM and parse its directive.

        Extraction strategy (ordered by reliability):
        1. XML <directive> tags (matches the proven <response> pattern from eb1.py)
        2. Bare JSON (first { to last })
        3. JSON inside markdown code fences
        4. Fallback to DIRECT on parse failure
        """
        response = self.conductor_model.response(
            context,
            system=system,
        )

        response_text = response.text if hasattr(response, "text") else str(response)

        # Strategy 1: XML <directive> tags (preferred)
        match = _DIRECTIVE_TAG_RE.search(response_text)
        if match:
            json_text = match.group(1).strip()
            try:
                data = json.loads(json_text)
                return ConductorDirective.model_validate(data)
            except (json.JSONDecodeError, ValueError) as e:
                _CONDUCTOR_LOG.warning(
                    "Conductor directive XML parse failed (trying bare JSON): %s",
                    e,
                )

        # Strategy 2: Bare JSON or code-fenced JSON
        json_text = response_text
        if "```json" in response_text:
            json_text = response_text.split("```json")[1].split("```")[0].strip()
        elif "```" in response_text:
            parts = response_text.split("```")
            if len(parts) >= 3:
                json_text = parts[1].strip()
        else:
            # Try to find bare JSON object
            start_idx = response_text.find("{")
            end_idx = response_text.rfind("}")
            if start_idx != -1 and end_idx > start_idx:
                json_text = response_text[start_idx : end_idx + 1]

        try:
            data = json.loads(json_text)
            return ConductorDirective.model_validate(data)
        except (json.JSONDecodeError, ValueError, KeyError) as e:
            _CONDUCTOR_LOG.warning(
                "All conductor directive parse strategies failed — "
                "falling back to DIRECT routing: %s. "
                "Raw response: %.500s",
                e,
                response_text,
            )
            return ConductorDirective(
                action=ConductorAction.DIRECT,
                reasoning=f"Parse error, falling back to direct: {e}",
                target_model=self.default_model,
            )

    def _execute_direct(
        self,
        query: str,
        system: str | None,
        target_model: str | None,
        trace: ConductorTrace,
    ) -> ChatResponse:
        """Execute a direct model call."""
        model_name = target_model or self.default_model
        model = self.model_pool.get(model_name, self.model_pool[self.default_model])

        start = time.perf_counter()
        response = model.response(query, system=system)
        latency = (time.perf_counter() - start) * 1000

        trace.total_latency_ms += latency
        return response.chat_response

    def _execute_subtask(
        self,
        subtask: SubtaskSpec,
        previous_results: dict[str, SubtaskResult],
        trace: ConductorTrace,
    ) -> SubtaskResult:
        """Execute a single subtask with GNN-informed routing."""
        # Build dependency context using XML tags (consistent with the tagging
        # pattern used in _update_context_with_results and _build_conductor_context).
        context_parts: list[str] = []
        for dep_id in subtask.depends_on:
            if dep_id in previous_results:
                dep_result = previous_results[dep_id]
                context_parts.append(
                    f"<<DEPENDENCY id={dep_id}>>\n"
                    f"{dep_result.response}\n"
                    f"<<END_DEPENDENCY>>"
                )

        full_query = subtask.query
        if context_parts:
            full_query = f"Context:\n{''.join(context_parts)}\n\nTask: {subtask.query}"

        # Determine model
        if subtask.preferred_model and subtask.preferred_model in self.model_pool:
            model_name = subtask.preferred_model
            signals = None
        else:
            # Use GNN to select
            signals = self.gnn_tool.get_routing_signals(full_query)
            model_name = signals.chosen_model

        model = self.model_pool.get(model_name, self.model_pool[self.default_model])

        # Execute
        start = time.perf_counter()
        response = model.response(full_query)
        latency = (time.perf_counter() - start) * 1000

        # Capture usage if the response carries it (ChatResponse.usage).
        usage = getattr(response, "usage", None)
        if usage is None:
            chat_resp = getattr(response, "chat_response", None)
            if chat_resp is not None:
                usage = getattr(chat_resp, "usage", None)

        result = SubtaskResult(
            subtask_id=subtask.id,
            model_used=model_name,
            response=response.text,
            latency_ms=latency,
            routing_signals=signals,
            usage=usage,
        )

        trace.subtask_results.append(result)
        trace.total_latency_ms += latency

        return result

    def _build_dependency_graph(
        self, subtasks: list[SubtaskSpec]
    ) -> dict[str, set[str]]:
        """Build dependency graph from subtask specifications.

        Args:
            subtasks: List of subtask specifications.

        Returns:
            Dictionary mapping subtask_id to set of dependency IDs.
        """
        graph: dict[str, set[str]] = {}
        for subtask in subtasks:
            graph[subtask.id] = set(subtask.depends_on)
        return graph

    def _validate_dependency_graph(
        self, graph: dict[str, set[str]], subtasks: list[SubtaskSpec]
    ) -> None:
        """Validate dependency graph has no cycles using Kahn's algorithm.

        Args:
            graph: Dependency graph (subtask_id -> dependencies).
            subtasks: List of subtask specifications.

        Raises:
            ValueError: If dependency graph contains a cycle.
        """
        # Compute in-degree for each node
        indegree: dict[str, int] = {st.id: len(graph[st.id]) for st in subtasks}

        # Start with nodes that have no dependencies
        queue: deque[str] = deque(
            [st_id for st_id, deg in indegree.items() if deg == 0]
        )
        processed_count = 0

        while queue:
            current = queue.popleft()
            processed_count += 1

            # Decrement in-degree of dependent nodes
            for st in subtasks:
                if current in graph[st.id]:
                    indegree[st.id] -= 1
                    if indegree[st.id] == 0:
                        queue.append(st.id)

        # If we didn't process all nodes, there's a cycle
        if processed_count != len(subtasks):
            remaining = [st.id for st in subtasks if indegree[st.id] > 0]
            raise ValueError(
                f"Dependency graph contains cycle. Remaining nodes: {remaining}. "
                f"Check depends_on fields for circular dependencies."
            )

    def _compute_execution_waves(
        self, graph: dict[str, set[str]], subtasks: list[SubtaskSpec]
    ) -> list[list[SubtaskSpec]]:
        """Compute execution waves for parallel subtask execution.

        Each wave contains subtasks that can execute in parallel (no dependencies
        between them). Waves are ordered such that all dependencies of subtasks
        in wave N are satisfied by waves 0..N-1.

        Args:
            graph: Dependency graph (subtask_id -> dependencies).
            subtasks: List of subtask specifications.

        Returns:
            List of waves, where each wave is a list of subtasks to execute in parallel.
        """
        subtask_map = {st.id: st for st in subtasks}
        indegree: dict[str, int] = {st.id: len(graph[st.id]) for st in subtasks}
        waves: list[list[SubtaskSpec]] = []

        while any(indegree[st.id] >= 0 for st in subtasks):
            # Find all subtasks with no remaining dependencies
            current_wave = [
                subtask_map[st_id]
                for st_id, deg in indegree.items()
                if deg == 0
            ]

            if not current_wave:
                break  # All remaining nodes processed

            waves.append(current_wave)

            # Mark current wave as processed
            for st in current_wave:
                indegree[st.id] = -1  # -1 indicates already processed

            # Decrement in-degree for dependent subtasks
            for st in subtasks:
                if indegree[st.id] >= 0:
                    for dep_id in graph[st.id]:
                        if indegree.get(dep_id) == -1:
                            indegree[st.id] -= 1

        return waves

    def _execute_subtasks_parallel(
        self,
        subtasks: list[SubtaskSpec],
        existing_outputs: dict[str, SubtaskResult],
        trace: ConductorTrace,
        budget_tracker: BudgetTracker,
    ) -> dict[str, SubtaskResult]:
        """Execute subtasks with dependency-aware parallelism.

        Algorithm:
        1. Build dependency graph from subtask.depends_on
        2. Validate graph (cycle detection)
        3. Compute topological waves for parallel execution
        4. Execute each wave with ThreadPoolExecutor
        5. Collect results, track metrics, handle errors

        Args:
            subtasks: List of subtask specifications to execute.
            existing_outputs: Existing results from previous iterations/subtasks.
            trace: Conductor trace for logging.
            budget_tracker: Budget tracker for enforcement.

        Returns:
            Dictionary mapping subtask_id to SubtaskResult.
        """
        # Step 1: Build dependency graph
        graph = self._build_dependency_graph(subtasks)

        # Step 2: Validate graph (cycle detection)
        self._validate_dependency_graph(graph, subtasks)

        # Step 3: Compute execution waves
        waves = self._compute_execution_waves(graph, subtasks)

        # Step 4: Execute waves with parallelism
        # Start with existing outputs so subtasks can depend on previous results
        results: dict[str, SubtaskResult] = dict(existing_outputs)
        parallelism_metrics = ParallelismMetrics()

        for wave_idx, wave_subtasks in enumerate(waves):
            budget_tracker.check_latency()

            wave_start = time.perf_counter()
            wave_results = self._execute_wave_parallel(
                wave_subtasks,
                results,  # Previous results for dependency resolution
                budget_tracker,
                parallelism_metrics,
            )
            wave_latency = (time.perf_counter() - wave_start) * 1000

            parallelism_metrics.execution_waves.append([st.id for st in wave_subtasks])
            parallelism_metrics.wave_latencies_ms.append(wave_latency)
            parallelism_metrics.total_parallel_latency_ms += wave_latency
            results.update(wave_results)

            for result in wave_results.values():
                trace.subtask_results.append(result)
                parallelism_metrics.total_sequential_latency_ms += result.latency_ms

        # Step 5: Compute critical path and speedup metrics
        self._compute_critical_path(parallelism_metrics, results, subtasks, graph)
        parallelism_metrics.compute_speedup()
        trace.parallelism_metrics = parallelism_metrics

        return results

    def _execute_wave_parallel(
        self,
        wave_subtasks: list[SubtaskSpec],
        previous_results: dict[str, SubtaskResult],
        budget_tracker: BudgetTracker,
        parallelism_metrics: ParallelismMetrics,
    ) -> dict[str, SubtaskResult]:
        """Execute a wave of independent subtasks in parallel.

        Args:
            wave_subtasks: Subtasks to execute (no dependencies between them).
            previous_results: Results from previous waves for dependency resolution.
            budget_tracker: Budget tracker for enforcement.
            parallelism_metrics: Metrics collector.

        Returns:
            Dictionary mapping subtask_id to SubtaskResult.
        """
        if len(wave_subtasks) == 1:
            # No parallelism benefit for single subtask
            subtask = wave_subtasks[0]
            try:
                budget_tracker.check_and_increment_subtask()
                # Create a temporary trace for this subtask
                temp_trace = ConductorTrace()
                result = self._execute_subtask(subtask, previous_results, temp_trace)
                parallelism_metrics.subtask_latencies_ms[subtask.id] = result.latency_ms
                return {subtask.id: result}
            except Exception as exc:
                _CONDUCTOR_LOG.warning(
                    "Subtask %s failed with %s: %s",
                    subtask.id,
                    type(exc).__name__,
                    exc,
                )
                parallelism_metrics.failed_subtasks.append(subtask.id)
                return {
                    subtask.id: SubtaskResult(
                        subtask_id=subtask.id,
                        model_used="error",
                        response=f"ERROR: {str(exc)}",
                        latency_ms=0.0,
                        error=exc,
                    )
                }

        # Parallel execution for multiple subtasks
        from ember.xcs.utils.executors import get_shared_executor

        executor = get_shared_executor(max_workers=len(wave_subtasks))

        # Submit all subtasks
        future_to_subtask = {}
        for subtask in wave_subtasks:
            try:
                budget_tracker.check_and_increment_subtask()
                temp_trace = ConductorTrace()
                future = executor.submit(
                    self._execute_subtask,
                    subtask,
                    previous_results,
                    temp_trace,
                )
                future_to_subtask[future] = subtask
            except BudgetExhaustedError as e:
                _CONDUCTOR_LOG.warning(
                    "Budget exhausted before submitting subtask %s "
                    "(dimension=%s, limit=%s, consumed=%s)",
                    subtask.id,
                    e.dimension,
                    e.limit,
                    e.consumed,
                )
                parallelism_metrics.failed_subtasks.append(subtask.id)

        # Collect results as they complete
        results: dict[str, SubtaskResult] = {}
        for future in as_completed(future_to_subtask):
            subtask = future_to_subtask[future]
            try:
                result = future.result()
                results[subtask.id] = result
                parallelism_metrics.subtask_latencies_ms[subtask.id] = result.latency_ms
            except Exception as exc:
                _CONDUCTOR_LOG.warning(
                    "Parallel subtask %s failed with %s: %s",
                    subtask.id,
                    type(exc).__name__,
                    exc,
                )
                parallelism_metrics.failed_subtasks.append(subtask.id)
                results[subtask.id] = SubtaskResult(
                    subtask_id=subtask.id,
                    model_used="error",
                    response=f"ERROR: {str(exc)}",
                    latency_ms=0.0,
                    error=exc,
                )

        return results

    def _compute_critical_path(
        self,
        parallelism_metrics: ParallelismMetrics,
        results: dict[str, SubtaskResult],
        subtasks: list[SubtaskSpec],
        graph: dict[str, set[str]],
    ) -> None:
        """Compute critical path (longest dependency chain) for speedup analysis.

        Args:
            parallelism_metrics: Metrics object to update.
            results: Completed subtask results.
            subtasks: All subtasks.
            graph: Dependency graph.
        """
        # Dynamic programming: compute longest path to each node
        subtask_map = {st.id: st for st in subtasks}
        dp: dict[str, float] = {}
        parent: dict[str, str | None] = {}

        def compute_longest_path(subtask_id: str) -> float:
            """Recursively compute longest path to this subtask."""
            if subtask_id in dp:
                return dp[subtask_id]

            result = results.get(subtask_id)
            if not result or result.error:
                dp[subtask_id] = 0.0
                parent[subtask_id] = None
                return 0.0

            # Base case: no dependencies
            if not graph[subtask_id]:
                dp[subtask_id] = result.latency_ms
                parent[subtask_id] = None
                return result.latency_ms

            # Recursive case: max(dependency paths) + current latency
            max_dep_path = 0.0
            max_dep_id = None
            for dep_id in graph[subtask_id]:
                dep_path = compute_longest_path(dep_id)
                if dep_path > max_dep_path:
                    max_dep_path = dep_path
                    max_dep_id = dep_id

            dp[subtask_id] = max_dep_path + result.latency_ms
            parent[subtask_id] = max_dep_id
            return dp[subtask_id]

        # Compute longest paths for all subtasks
        for subtask in subtasks:
            compute_longest_path(subtask.id)

        # Find the overall critical path (maximum among all endpoints)
        if dp:
            critical_end = max(dp.items(), key=lambda x: x[1])
            parallelism_metrics.critical_path_latency_ms = critical_end[1]

            # Reconstruct critical path
            path = []
            current = critical_end[0]
            while current is not None:
                path.append(current)
                current = parent.get(current)
            parallelism_metrics.critical_path = list(reversed(path))

    def _synthesize_results(
        self,
        original_query: str,
        system: str | None,
        subtask_results: dict[str, SubtaskResult],
        synthesis_prompt: str | None,
        trace: ConductorTrace,
    ) -> ChatResponse:
        """Synthesize subtask results into final answer."""
        # Build synthesis context
        results_text = "\n\n".join(
            [f"## {result.subtask_id}\n{result.response}" for result in subtask_results.values()]
        )

        default_synthesis = f"""Original Question: {original_query}

Subtask Results:
{results_text}

Please synthesize these results into a comprehensive, coherent answer to the original question."""

        synthesis_query = synthesis_prompt or default_synthesis

        # Use the conductor model for synthesis
        start = time.perf_counter()
        response = self.conductor_model.response(synthesis_query, system=system)
        latency = (time.perf_counter() - start) * 1000

        trace.total_latency_ms += latency
        return response.chat_response

    def _synthesize_partial(
        self,
        original_query: str,
        system: str | None,
        subtask_results: dict[str, SubtaskResult],
        trace: ConductorTrace,
    ) -> ChatResponse:
        """Synthesize partial results when budget is exhausted.

        Provides best-effort answer using completed subtasks only.
        If no subtasks completed, falls back to top-tier model.

        Args:
            original_query: The user's original query.
            system: Optional system prompt.
            subtask_results: Completed subtask results (may be empty).
            trace: Execution trace for metrics.

        Returns:
            Best-effort final response.
        """
        # If we have some results, synthesize them
        if subtask_results:
            # Filter out failed subtasks
            successful_results = {
                k: v for k, v in subtask_results.items() if v.error is None
            }

            if successful_results:
                results_text = "\n\n".join(
                    [
                        f"## Partial Result {i+1} ({result.subtask_id})\n{result.response}"
                        for i, result in enumerate(successful_results.values())
                    ]
                )

                partial_synthesis = f"""Original Question: {original_query}

NOTE: Budget limits prevented completing all subtasks. Here are the partial results:

{results_text}

Based on these partial results, provide the best answer you can to the original question.
Acknowledge any limitations due to incomplete analysis."""

                start = time.perf_counter()
                response = self.conductor_model.response(partial_synthesis, system=system)
                latency = (time.perf_counter() - start) * 1000
                trace.total_latency_ms += latency
                return response.chat_response

        # No results - fall back to top-tier model
        return self._fallback_response(original_query, system, trace)

    def _delegate_to_eb1(
        self,
        query: str,
        system: str | None,
        config: str | None,
        trace: ConductorTrace,
    ) -> ChatResponse:
        """Delegate to full EB1 ensemble."""
        # Import here to avoid circular dependency
        from .eb1 import (
            DEFAULT_PLAN,
            DEFAULT_PRO_PLAN,
            Eb1Graph,
            PromptContext,
        )

        plan = DEFAULT_PRO_PLAN if config == "eb1-pro" else DEFAULT_PLAN
        runner = Eb1Graph(plan)
        context = PromptContext(prompt=query, system=system)

        start = time.perf_counter()
        result = runner(context)
        latency = (time.perf_counter() - start) * 1000

        trace.total_latency_ms += latency
        return result

    def _update_context_with_results(
        self,
        context: str,
        directive: ConductorDirective,
        results: dict[str, SubtaskResult],
    ) -> str:
        """Update context with subtask execution results."""
        results_section = "\n".join(
            [
                f"<<SUBTASK_RESULT id={r.subtask_id} model={r.model_used}>>\n"
                f"{r.response}\n<<END_SUBTASK_RESULT>>"
                for r in results.values()
            ]
        )
        return f"{context}\n\n{results_section}"

    def _fallback_response(
        self,
        query: str,
        system: str | None,
        trace: ConductorTrace,
    ) -> ChatResponse:
        """Fallback when max iterations reached."""
        model = self.model_pool[self.default_model]

        start = time.perf_counter()
        response = model.response(query, system=system)
        latency = (time.perf_counter() - start) * 1000

        trace.total_latency_ms += latency
        return response.chat_response


# =============================================================================
# Top-Level EB1-GNN Inference Router
# =============================================================================


@dataclass
class EB1GNNResult:
    """Result from EB1-GNN inference."""

    response: ChatResponse
    routing_path: str  # "trivial", "easy", "hard"
    routing_signals: RoutingSignals
    model_used: str | list[str]  # Single model or list if conductor
    latency_ms: float
    conductor_trace: ConductorTrace | None = None
    gnn_call_record: GNNCallRecord | None = None  # Record of initial GNN call


class EB1GNNInference(Operator):
    """Top-level inference router combining GNN and LLM conductor.

    This is the main entry point for EB1-GNN inference. It:
    1. Computes routing signals from the GNN
    2. Classifies the query into a tier (trivial, easy, hard)
    3. Routes to appropriate execution path

    Routing Paths:
        TRIVIAL -> Cheap model (Gemini Flash, GPT-5 Nano)
        EASY -> GNN's top-1 prediction (direct call)
        HARD -> EB1 Conductor (LLM orchestration with GNN tool)

    Example:
        inference = EB1GNNInference.from_registry("eb1-gnn-preview-12-19")

        # Simple query - fast path
        result = inference("What is 2+2?")
        assert result.routing_path == "trivial"

        # Complex query - conductor path
        result = inference("Compare and contrast the economic policies of...")
        assert result.routing_path == "conductor"
        print(result.conductor_trace.subtask_results)
    """

    gnn_tool: GNNRoutingTool
    conductor: EB1Conductor
    model_pool: dict[str, ModelBinding]
    default_model: str
    thresholds: RoutingThresholds
    cheap_model_name: str
    _metrics: ConductorMetrics

    def __init__(
        self,
        gnn_router: GNNRouterProtocol,
        thresholds: RoutingThresholds | None = None,
        conductor_model: str = CLAUDE_OPUS,
        cheap_model: str = "gemini-2.5-flash-lite",  # Cheapest in GNN 8-LLM set
        model_pool: dict[str, str] | None = None,
        enable_metrics: bool = True,
        default_model: str = "Claude Opus 4.5",
    ):
        """Initialize the inference router.

        Args:
            gnn_router: Any object implementing GNNRouterProtocol.
            thresholds: Custom routing thresholds.
            conductor_model: Model for the conductor LLM.
            cheap_model: Model ID for trivial queries.
            model_pool: Available models for execution.
            enable_metrics: Whether to track detailed metrics (default True).
            default_model: Display name of the default/strongest model in pool.
        """
        self.gnn_tool = GNNRoutingTool(gnn_router)
        self.thresholds = thresholds or RoutingThresholds()
        self._enable_metrics = enable_metrics
        self._metrics = ConductorMetrics()

        pool_config = model_pool or _DEFAULT_MODEL_POOL
        self.model_pool = {
            name: models.instance(model_id) for name, model_id in pool_config.items()
        }

        self.default_model = default_model if default_model in pool_config else next(iter(pool_config))
        self.cheap_model_name = cheap_model
        self.conductor = EB1Conductor(
            gnn_router=gnn_router,
            conductor_model=conductor_model,
            model_pool=pool_config,
            default_model=default_model,
        )

    @property
    def metrics(self) -> ConductorMetrics:
        """Get the current metrics. Use reset_metrics() to clear."""
        return self._metrics

    def reset_metrics(self) -> None:
        """Reset all metrics to zero. Call before a new evaluation run."""
        self._metrics = ConductorMetrics()

    @classmethod
    def from_registry(
        cls,
        registry_name: str,
        registry_root: str | Path | None = None,
        device: str = "cpu",
        **kwargs: Any,  # Forwarded to __init__ (thresholds, conductor_model, etc.)
    ) -> "EB1GNNInference":
        """Load from model registry.

        Args:
            registry_name: Registry entry name (e.g., "eb1-gnn-preview-12-19").
            registry_root: Root directory for registry.
            device: Device for GNN inference.
            **kwargs: Additional arguments for __init__.

        Returns:
            Initialized EB1GNNInference instance.
        """
        # Import here to avoid circular dependency
        try:
            from embergraphrouter.inference.eb1_router import EB1Router
        except ImportError:
            from embergraphrouter.inference import EB1Router

        gnn_router = EB1Router.from_registry(
            registry_name,
            registry_root=registry_root,
            device=device,
        )

        return cls(gnn_router=gnn_router, **kwargs)

    def forward(
        self,
        query: str,
        *,
        system: str | None = None,
        force_path: str | None = None,
    ) -> EB1GNNResult:
        """Route and execute a query.

        Args:
            query: The user's query.
            system: Optional system prompt.
            force_path: Force a specific path ("trivial", "easy", "hard").

        Returns:
            EB1GNNResult with response and routing metadata.
        """
        start_time = time.perf_counter()
        gnn_start = time.perf_counter()

        # Get routing signals from GNN
        signals = self.gnn_tool.get_routing_signals(query)
        gnn_latency = (time.perf_counter() - gnn_start) * 1000

        # Determine routing path
        if force_path:
            path = force_path
            tier = RoutingTier(path) if path in [t.value for t in RoutingTier] else RoutingTier.HARD
        else:
            tier = signals.classify_with_thresholds(self.thresholds)
            path = tier.value

        # Record GNN call for metrics
        gnn_record = None
        if self._enable_metrics:
            gnn_record = GNNCallRecord.from_signals(
                query=query,
                signals=signals,
                latency_ms=gnn_latency,
                caller="initial",
            )
            self._metrics.record_gnn_call(gnn_record)
            self._metrics.record_path(tier)

        # Execute based on path
        exec_start = time.perf_counter()
        if path == "trivial":
            response, model_used = self._execute_trivial(query, system)
            conductor_trace = None
        elif path == "easy":
            response, model_used = self._execute_fast(query, system, signals)
            conductor_trace = None
        else:  # hard -> conductor
            response, conductor_trace = self.conductor(
                query, system=system, routing_signals=signals
            )
            # Extract models used from trace
            model_used = [r.model_used for r in conductor_trace.subtask_results]
            if not model_used:
                model_used = ["conductor"]

            # Record conductor metrics
            if self._enable_metrics and conductor_trace:
                for directive in conductor_trace.directives:
                    self._metrics.record_action(directive.action.value)
                for subtask_result in conductor_trace.subtask_results:
                    used_gnn = subtask_result.routing_signals is not None
                    self._metrics.record_subtask(used_gnn, had_override=False)
                self._metrics.conductor_latency_ms += conductor_trace.total_latency_ms

        exec_latency = (time.perf_counter() - exec_start) * 1000
        total_latency = (time.perf_counter() - start_time) * 1000

        # Record model usage
        if self._enable_metrics:
            if isinstance(model_used, str):
                self._metrics.record_model_usage(model_used)
            else:
                for model in model_used:
                    self._metrics.record_model_usage(model)
            self._metrics.execution_latency_ms += exec_latency
            self._metrics.total_latency_ms += total_latency

        return EB1GNNResult(
            response=response,
            routing_path=path,
            routing_signals=signals,
            model_used=model_used,
            latency_ms=total_latency,
            conductor_trace=conductor_trace,
            gnn_call_record=gnn_record,
        )

    def _execute_trivial(
        self,
        query: str,
        system: str | None,
    ) -> tuple[ChatResponse, str]:
        """Execute trivial query with cheap model."""
        # Find cheap model in pool, or use the configured one
        cheap_model = None
        for name, binding in self.model_pool.items():
            if "nano" in name.lower() or "flash" in name.lower():
                cheap_model = binding
                model_name = name
                break

        if cheap_model is None:
            cheap_model = models.instance(self.cheap_model_name)
            model_name = self.cheap_model_name

        response = cheap_model.response(query, system=system)
        return response.chat_response, model_name

    def _execute_fast(
        self,
        query: str,
        system: str | None,
        signals: RoutingSignals,
    ) -> tuple[ChatResponse, str]:
        """Execute with GNN's top-1 prediction."""
        model_name = signals.chosen_model
        model = self.model_pool.get(model_name)

        if model is None:
            model_name = self.default_model
            model = self.model_pool[model_name]

        response = model.response(query, system=system)
        return response.chat_response, model_name


# =============================================================================
# Exports
# =============================================================================

__all__ = [
    # Protocols
    "GNNRouteResult",
    "GNNRouterProtocol",
    # Routing core
    "RoutingTier",
    "RoutingThresholds",
    "RoutingSignals",
    "GNNRoutingTool",
    # Metrics
    "GNNCallRecord",
    "ConductorMetrics",
    # Conductor schemas
    "ConductorAction",
    "SubtaskSpec",
    "ConductorDirective",
    "SubtaskResult",
    "ConductorTrace",
    # Main classes
    "EB1Conductor",
    "EB1GNNResult",
    "EB1GNNInference",
]
