"""Data models for BTQL Agent."""

from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Callable


class QueryStatus(Enum):
    """Status of a query execution."""
    SUCCESS = "success"
    ERROR = "error"
    PARTIAL = "partial"


@dataclass
class QueryResult:
    """Result of a BTQL query execution."""
    status: QueryStatus
    data: list[dict[str, Any]] | None = None
    error: str | None = None
    raw_response: dict[str, Any] | None = None

    @property
    def is_success(self) -> bool:
        return self.status == QueryStatus.SUCCESS


@dataclass
class AgentConfig:
    """Configuration for the BTQL Agent.

    Supports four model client types:
    - "azure": Azure OpenAI (requires azure_openai_endpoint, azure_openai_deployment, token_provider)
    - "papyrus": Papyrus (requires papyrus_endpoint, papyrus_deployment, token_provider)
    - "copilot": Copilot via VS Code extension proxy (requires copilot_model)
    - "local": Local OpenAI-compatible server (requires local_model)
    """
    braintrust_api_key: str
    braintrust_org: str = "non-pme"  # "bt-azure" or "non-pme"
    max_iterations: int = 5
    timeout: int = 60
    default_filters: list[str] = field(default_factory=list)

    # Model client configuration
    model_client_type: str = "azure"  # "azure" | "papyrus" | "copilot" | "local"

    # Azure-specific (used when model_client_type == "azure")
    azure_openai_endpoint: str | None = None
    azure_openai_deployment: str | None = None
    token_provider: Callable[[], str] | None = None  # Azure AD token provider (shared by Azure and Papyrus)

    # Papyrus-specific (used when model_client_type == "papyrus")
    papyrus_endpoint: str | None = None
    papyrus_deployment: str | None = None

    # Copilot-specific (used when model_client_type == "copilot")
    copilot_endpoint: str = "http://localhost:8237"  # Backend server URL
    copilot_model: str | None = None  # Model ID from Copilot

    # Local-specific (used when model_client_type == "local")
    local_endpoint: str = "http://localhost:23333/api/openai/v1"  # Local OpenAI-compatible base URL
    local_model: str | None = None  # Local model name


@dataclass
class JudgeResponse:
    """Judge response explaining why a metric got its score."""
    metric: str
    score: float
    explanation: str
    root_span_id: str | None = None


@dataclass
class AgentResult:
    """Final result from the BTQL Agent."""
    success: bool
    metrics: dict[str, float] | None = None
    judge_responses: list[JudgeResponse] | None = None
    final_query: str | None = None
    iterations: int = 0
    error: str | None = None
    history: list[dict[str, Any]] = field(default_factory=list)


@dataclass
class RefinementAttempt:
    """Record of a single refinement attempt."""
    iteration: int
    query: str
    result: QueryResult
    refinement_reason: str | None = None


@dataclass
class SchemaInfo:
    """Discovered schema information for an experiment.

    Contains paths to where scores and explanations are stored in the data.
    Paths use dot notation (e.g., "scores.completeness", "metadata.completeness.explanation").
    """
    scores_path: str = "scores"  # Base path to scores object
    available_metrics: list[str] = field(default_factory=list)  # Discovered metric names
    metric_types: dict[str, str] = field(default_factory=dict)  # metric -> llm_judged|rule_based|hybrid
    metric_groups: list[dict[str, Any]] = field(default_factory=list)  # scorer-row metric groups from sampled rows
    explanation_paths: list[str] = field(default_factory=list)  # Paths to explanation fields
    id_field: str = "root_span_id"  # Field used to identify rows
    raw_sample: dict[str, Any] | None = None  # Sample row for reference
    # Cached query templates (with {experiment_id} and {root_span_id} as placeholders)
    scores_query_template: str | None = None
    judge_query_template: str | None = None
    trace_query_template: str | None = None
    # Grouped templates (one per scorer-like metric group)
    judge_query_templates: list[dict[str, Any]] = field(default_factory=list)
    # Internal-only template used during discovery verification (not persisted)
    explanation_query_template: str | None = None  # Temporary query for testing explanation extraction
