"""Execution trace domain entities for Yagra Run & Observe (G-14, G-15)."""

from __future__ import annotations

import importlib.metadata
import platform
from datetime import datetime
from enum import StrEnum
from typing import Any
from uuid import UUID, uuid4

from pydantic import BaseModel, ConfigDict, Field


class NodeStatus(StrEnum):
    """Execution status of a single node."""

    SUCCESS = "success"
    ERROR = "error"
    SKIPPED = "skipped"


class LLMCallTrace(BaseModel):
    """Token usage and cost information for a single LLM invocation.

    Populated only when the node uses a built-in LLM handler
    (llm, structured_llm, or streaming_llm).

    Attributes:
        model: Full litellm model string (e.g. 'openai/gpt-4o-mini').
        provider: Provider portion extracted from model (e.g. 'openai').
        prompt_tokens: Number of input tokens consumed.
        completion_tokens: Number of output tokens generated.
        total_tokens: Sum of prompt and completion tokens.
        estimated_cost_usd: Estimated cost in USD based on the pricing table.
            None if the model is not in the pricing table.
    """

    model_config = ConfigDict(extra="forbid")

    model: str = Field(
        description="Full litellm model string, e.g. 'openai/gpt-4o-mini'.",
        examples=["openai/gpt-4o-mini", "anthropic/claude-3-5-sonnet-20241022"],
    )
    provider: str = Field(
        description="LLM provider extracted from model, e.g. 'openai'.",
        examples=["openai", "anthropic", "google"],
    )
    prompt_tokens: int = Field(
        ge=0,
        description="Number of input tokens consumed by this LLM call.",
    )
    completion_tokens: int = Field(
        ge=0,
        description="Number of output tokens generated by this LLM call.",
    )
    total_tokens: int = Field(
        ge=0,
        description="Total tokens (prompt + completion).",
    )
    estimated_cost_usd: float | None = Field(
        default=None,
        description=(
            "Estimated cost in USD based on a static pricing table. "
            "None if the model is not in the table."
        ),
    )


class ErrorTrace(BaseModel):
    """Error information captured when a node raises an exception.

    Attributes:
        error_type: Fully qualified exception class name.
        error_message: str(exception) truncated to 2000 chars.
        traceback: Full traceback string. May be omitted in production environments.
    """

    model_config = ConfigDict(extra="forbid")

    error_type: str = Field(
        description="Fully qualified exception class name.",
        examples=["yagra.handlers.llm_handler.LLMHandlerCallError", "ValueError"],
    )
    error_message: str = Field(
        description="String representation of the exception, truncated to 2000 characters.",
    )
    traceback: str | None = Field(
        default=None,
        description="Full traceback string. Set to None when include_traceback=False.",
    )


class NodeTrace(BaseModel):
    """Execution record for a single workflow node.

    Attributes:
        node_id: Node identifier as declared in the workflow YAML.
        handler: Handler name used by the node.
        status: Execution result status.
        started_at: UTC ISO-8601 timestamp when the node started.
        ended_at: UTC ISO-8601 timestamp when the node ended.
        duration_ms: Wall-clock execution time in milliseconds.
        input_snapshot: Shallow copy of state keys read by the node.
        output_snapshot: State keys written by the node.
        llm_call: LLM token and cost details. None for non-LLM handlers.
        error: Exception details. None when status is SUCCESS.
    """

    model_config = ConfigDict(extra="forbid")

    node_id: str = Field(
        description="Node identifier as declared in the workflow YAML.",
        examples=["translate", "summarize", "classify"],
    )
    handler: str = Field(
        description="Handler name resolved at runtime.",
        examples=["llm", "structured_llm", "my_custom_handler"],
    )
    status: NodeStatus = Field(
        description="Execution result: 'success', 'error', or 'skipped'.",
    )
    started_at: datetime = Field(
        description="UTC timestamp when the node execution started.",
    )
    ended_at: datetime = Field(
        description="UTC timestamp when the node execution ended.",
    )
    duration_ms: float = Field(
        ge=0.0,
        description="Wall-clock execution time in milliseconds.",
    )
    input_snapshot: dict[str, Any] = Field(
        default_factory=dict,
        description=(
            "Shallow copy of state keys read by the node before execution. "
            "Non-JSON-serializable values are replaced with their repr() string."
        ),
    )
    output_snapshot: dict[str, Any] = Field(
        default_factory=dict,
        description=(
            "State keys returned by the node handler. "
            "Non-JSON-serializable values are replaced with their repr() string."
        ),
    )
    llm_call: LLMCallTrace | None = Field(
        default=None,
        description="LLM token usage and cost. None for non-LLM handlers.",
    )
    error: ErrorTrace | None = Field(
        default=None,
        description="Exception details. None when status is 'success'.",
    )


class RunSummary(BaseModel):
    """Aggregate statistics for a single workflow run.

    Attributes:
        total_nodes_executed: Count of nodes that were actually executed.
        succeeded_nodes: Count of nodes with status SUCCESS.
        failed_nodes: Count of nodes with status ERROR.
        total_duration_ms: End-to-end wall-clock time in milliseconds.
        total_prompt_tokens: Aggregated prompt tokens across all LLM calls.
        total_completion_tokens: Aggregated completion tokens across all LLM calls.
        total_tokens: Sum of all token counts.
        total_estimated_cost_usd: Aggregated estimated cost in USD.
        node_order: Ordered list of node IDs in execution sequence.
    """

    model_config = ConfigDict(extra="forbid")

    total_nodes_executed: int = Field(ge=0)
    succeeded_nodes: int = Field(ge=0)
    failed_nodes: int = Field(ge=0)
    total_duration_ms: float = Field(ge=0.0)
    total_prompt_tokens: int = Field(ge=0, default=0)
    total_completion_tokens: int = Field(ge=0, default=0)
    total_tokens: int = Field(ge=0, default=0)
    total_estimated_cost_usd: float | None = Field(
        default=None,
        description="Aggregated estimated cost in USD. None if no pricing data was available.",
    )
    node_order: list[str] = Field(
        default_factory=list,
        description="Execution sequence of node IDs, in order of completion.",
    )


class WorkflowRunTrace(BaseModel):
    """Complete execution trace for a single Yagra.invoke() call.

    This is the top-level document written to a trace file.
    One WorkflowRunTrace is produced per Yagra.invoke() call when tracing is enabled.

    Attributes:
        schema_version: Trace format version for future compatibility (currently '1.0').
        run_id: UUID v4 uniquely identifying this execution run.
        workflow_name: Stem of the workflow YAML file.
        workflow_version: Version string from the workflow YAML (GraphSpec.version).
        workflow_path: Absolute path to the workflow YAML file, if available.
        started_at: UTC timestamp when Yagra.invoke() was called.
        ended_at: UTC timestamp when Yagra.invoke() returned (or raised).
        status: Overall run status. 'error' if any node failed.
        nodes: Ordered list of NodeTrace records.
        summary: Aggregated statistics for the entire run.
        metadata: Environment and version information for reproducibility.
    """

    model_config = ConfigDict(extra="forbid")

    schema_version: str = Field(
        default="1.0",
        description="Trace format schema version. Currently '1.0'.",
    )
    run_id: UUID = Field(
        default_factory=uuid4,
        description="UUID v4 uniquely identifying this run.",
    )
    workflow_name: str = Field(
        description="Stem of the workflow YAML file.",
        examples=["translate", "multi-agent-pipeline"],
    )
    workflow_version: str = Field(
        description="Version field from the workflow YAML (GraphSpec.version).",
        examples=["1.0"],
    )
    workflow_path: str | None = Field(
        default=None,
        description="Absolute path to the workflow YAML file. None if not available.",
    )
    started_at: datetime = Field(
        description="UTC timestamp when Yagra.invoke() was called.",
    )
    ended_at: datetime = Field(
        description="UTC timestamp when Yagra.invoke() returned or raised.",
    )
    status: NodeStatus = Field(
        description="Overall run status. 'error' if any node in nodes[] has status 'error'.",
    )
    nodes: list[NodeTrace] = Field(
        default_factory=list,
        description="Ordered list of node execution records.",
    )
    summary: RunSummary = Field(
        description="Aggregated statistics computed from nodes[].",
    )
    metadata: dict[str, str] = Field(
        default_factory=dict,
        description=(
            "Environment metadata: yagra_version, python_version, platform. "
            "Useful for reproducibility analysis."
        ),
    )

    @classmethod
    def build_metadata(cls) -> dict[str, str]:
        """Collects runtime environment metadata.

        Returns:
            Dict with 'yagra_version', 'python_version', and 'platform'.
        """
        try:
            yagra_version = importlib.metadata.version("yagra")
        except importlib.metadata.PackageNotFoundError:
            yagra_version = "unknown"
        return {
            "yagra_version": yagra_version,
            "python_version": platform.python_version(),
            "platform": platform.system(),
        }

    @classmethod
    def compute_summary(cls, nodes: list[NodeTrace], total_duration_ms: float) -> RunSummary:
        """Computes RunSummary from a completed list of NodeTrace records.

        Args:
            nodes: List of completed node traces.
            total_duration_ms: Total wall-clock duration of the entire run.

        Returns:
            RunSummary with aggregated statistics.
        """
        succeeded = sum(1 for n in nodes if n.status == NodeStatus.SUCCESS)
        failed = sum(1 for n in nodes if n.status == NodeStatus.ERROR)

        prompt_tokens = 0
        completion_tokens = 0
        total_tokens = 0
        estimated_cost: float | None = None

        for node in nodes:
            if node.llm_call is not None:
                prompt_tokens += node.llm_call.prompt_tokens
                completion_tokens += node.llm_call.completion_tokens
                total_tokens += node.llm_call.total_tokens
                if node.llm_call.estimated_cost_usd is not None:
                    estimated_cost = (estimated_cost or 0.0) + node.llm_call.estimated_cost_usd

        return RunSummary(
            total_nodes_executed=len(nodes),
            succeeded_nodes=succeeded,
            failed_nodes=failed,
            total_duration_ms=total_duration_ms,
            total_prompt_tokens=prompt_tokens,
            total_completion_tokens=completion_tokens,
            total_tokens=total_tokens,
            total_estimated_cost_usd=estimated_cost,
            node_order=[n.node_id for n in nodes],
        )
