"""
Context Enricher — builds EnrichedContext from all available SDK components.

Purely local (no network calls). Assembles data from:
1. ContextAggregator — previous spans in this trace
2. PatternCache — workflow/customer data synced from platform

Target: <2ms (all local dict lookups).
"""

import logging
from typing import TYPE_CHECKING, Any

from .models import EnrichedContext

if TYPE_CHECKING:
    from .context_aggregator import ContextAggregator

logger = logging.getLogger("aigie.judge.context_enricher")


class ContextEnricher:
    """Builds EnrichedContext from all available SDK components."""

    def __init__(
        self,
        context_aggregator: "ContextAggregator",
        pattern_cache: Any = None,
    ):
        self._ctx_agg = context_aggregator
        self._pattern_cache = pattern_cache

    def enrich(
        self,
        trace_id: str,
        span_id: str,
        span_type: str,
        signal_hints: list[str],
        metadata: dict[str, Any] | None = None,
    ) -> EnrichedContext:
        """Build enriched context. Target: <2ms (all local data)."""
        metadata = metadata or {}

        # Step position from aggregator
        step_index = 0
        total_steps = 0
        previous_step_had_issue = False
        previous_step_issue_type = None
        consecutive_issues = 0
        span_name = metadata.get("name", "")

        agg_ctx = self._ctx_agg.get_context(trace_id)
        if agg_ctx and agg_ctx.spans:
            completed = [s for s in agg_ctx.spans if s.status in ("success", "error")]
            total_steps = len(completed)
            step_index = total_steps  # Current span is the next step

            # Check previous step
            if completed:
                last = completed[-1]
                if last.issues_detected:
                    previous_step_had_issue = True
                    previous_step_issue_type = (
                        last.issues_detected[0] if last.issues_detected else None
                    )

                # Count consecutive issues (walk backwards)
                for s in reversed(completed):
                    if s.issues_detected:
                        consecutive_issues += 1
                    else:
                        break

            # Get span name from aggregator if available
            for s in agg_ctx.spans:
                if s.span_id == span_id:
                    span_name = s.name or span_name
                    break

        # Workflow/customer data from PatternCache
        workflow_failure_points: list[str] = []
        workflow_common_issues: list[str] = []
        customer_top_issues: list[str] = []
        judge_accuracy_scores: dict[str, float] = {}

        if self._pattern_cache and hasattr(self._pattern_cache, "get_enrichment_context"):
            enrichment = self._pattern_cache.get_enrichment_context(
                agent_name=metadata.get("agent_name"),
                workflow_id=metadata.get("workflow_id"),
            )
            if enrichment:
                workflow_failure_points = enrichment.get("workflow_failure_points", [])
                workflow_common_issues = enrichment.get("workflow_common_issues", [])
                customer_top_issues = enrichment.get("customer_top_issues", [])
                judge_accuracy_scores = enrichment.get("judge_accuracy", {})

        return EnrichedContext(
            signal_hints=signal_hints,
            span_type=span_type,
            span_name=span_name,
            model=metadata.get("model"),
            tool_name=metadata.get("tool_name"),
            step_index=step_index,
            total_steps_so_far=total_steps,
            previous_step_had_issue=previous_step_had_issue,
            previous_step_issue_type=previous_step_issue_type,
            consecutive_issues=consecutive_issues,
            workflow_id=metadata.get("workflow_id"),
            workflow_name=metadata.get("workflow_name"),
            workflow_known_failure_points=workflow_failure_points,
            workflow_common_issues=workflow_common_issues,
            agent_name=metadata.get("agent_name"),
            agent_framework=metadata.get("agent_framework"),
            customer_top_issues=customer_top_issues,
            judge_accuracy_scores=judge_accuracy_scores,
        )
