from __future__ import annotations

from datetime import datetime
from enum import Enum
from typing import Any

from pydantic import BaseModel, Field


class User(BaseModel):
    id: str
    created_at: datetime = Field(default_factory=datetime.utcnow)
    last_login: datetime = Field(default_factory=datetime.utcnow)


class SessionStatus(str, Enum):
    IDLE = "idle"
    RUNNING = "running"
    PAUSED = "paused"
    COMPLETED = "completed"
    ERROR = "error"


class Session(BaseModel):
    id: str
    user_id: str  # Owner of this session
    name: str
    status: SessionStatus = SessionStatus.IDLE
    task_type: str
    created_at: datetime = Field(default_factory=datetime.utcnow)
    updated_at: datetime = Field(default_factory=datetime.utcnow)
    current_round: int = 0
    max_rounds: int = 20


class LiquidPromptEntry(BaseModel):
    """Represents a single liquid prompt file path with analysis flag."""
    path: str
    analyze_structure: bool = True  # Whether to analyze the structure of this file


class MetricConfig(BaseModel):
    """Configuration for a single metric."""
    name: str
    is_llm_judged: bool = True  # True for LLM-judged metrics, False for rule-based metrics
    description: str | None = None  # Description of what this metric measures
    measurement_purpose: str | None = None  # Purpose of this measurement
    optimization_direction: str = "higher_better"  # "higher_better" or "lower_better"


class SessionConfig(BaseModel):
    task_type: str
    goal: str | None = None
    dataset_name: str
    metrics_config: list[MetricConfig] | None = None  # Detailed configuration for each metric (replaces key_metrics)
    picasso_version: str | None = None
    conversation_mode: str | None = None
    max_dataset_rows: int | None = None
    dataset_sampling_method: str | None = None  # "all", "first", "last", or "custom"
    additional_cmd_args: str | None = None       # Additional arguments appended to eval runner command
    eval_project_name: str | None = None
    dataset_project_name: str | None = None
    branch: str
    original_prompt_file_path: str | None = None
    liquid_prompt_file_path: str | None = None  # Deprecated, kept for backward compatibility
    liquid_prompt_entries: list[LiquidPromptEntry] | None = None  # New: multiple entries with analysis flags
    termination_rule: dict[str, Any] | None = None
    max_rounds: int = 20
    commit_id: str | None = None
    eval_name: str | None = None
    add_feature_flags: str | None = None
    baseline_feature_flags: str | None = None
    model_client_type: str | None = "Azure"  # "Azure", "Papyrus", "Copilot", or "Local"
    copilot_model: str | None = None         # Copilot model ID (from VS Code extension)
    model_deployment: str | None = None      # Azure OpenAI deployment name
    azure_endpoint: str | None = None         # Azure OpenAI endpoint URL
    azure_api_version: str | None = None     # Azure OpenAI API version
    papyrus_model: str | None = None         # Papyrus model name
    papyrus_endpoint: str | None = None      # Papyrus endpoint URL
    papyrus_quota_id: str | None = None       # Papyrus quota ID
    papyrus_scope: str | None = None           # Papyrus auth scope for token generation
    local_model: str | None = None           # Local model ID (e.g., "gpt-5.1-codex")
    local_endpoint: str | None = None        # Local OpenAI-compatible base URL (default: http://localhost:23333/api/openai/v1)
    selection_strategy: str | None = "greedy"  # "greedy" or "gepa"
    pareto_max_selected: int | None = Field(default=2, ge=1, le=5)  # Max candidates from Pareto front
    baseline_experiment_id: str | None = None    # Skip TestExecutor if provided
    baseline_experiment_link: str | None = None  # Braintrust experiment URL
    local_debug: bool = False                # Use mock data instead of real evals
    badcase_deep_analysis: bool = False      # Enable deeper per-case bad-case analysis
    run_leaderboard_after_optimization: bool = False  # Trigger leaderboard eval with best branch after optimization
    braintrust_org: str | None = "non-pme"   # "bt-azure" or "non-pme"
    environment: str | None = None           # "china" or "global"
    model: str | None = None                 # "default", "gpt4", "gpt5", "gpt5.1", or "gpt5.2"
    picasso_env: str | None = None           # "runner" or "custom" — maps to -PicassoEnv in PS1 script
    task_id_prefix: str | None = None        # Optional custom task ID prefix


class Checkpoint(BaseModel):
    path: str
    agent_turn: int
    agent_name: str
    saved_at: str


class AgentEventType(str, Enum):
    AGENT_START = "agent_start"
    AGENT_MESSAGE = "agent_message"
    AGENT_COMPLETE = "agent_complete"
    ROUTING = "routing"
    TOOL_CALL = "tool_call"


class AgentEvent(BaseModel):
    type: AgentEventType
    timestamp: datetime = Field(default_factory=datetime.utcnow)
    agent: str
    payload: dict[str, Any] = Field(default_factory=dict)


class MetricsSnapshot(BaseModel):
    version: int
    timestamp: datetime
    overall_score: float
    dimensions: dict[str, float] = Field(default_factory=dict)
    experiment_id: str
    experiment_link: str


class BadCase(BaseModel):
    root_span_id: str
    experiment_id: str
    test_model_response: str = ""
    judge_response: str = ""
    judge_reasoning: str = ""
    merged_scores: dict[str, float] = Field(default_factory=dict)
    average_score: float = 0.0
    low_dimensions: list[str] = Field(default_factory=list)


class BadCaseFilters(BaseModel):
    experiment_id: str | None = None
    min_score: float | None = None
    max_score: float | None = None
    dimensions: list[str] | None = None
    keyword: str | None = None


class PromptFile(BaseModel):
    name: str
    path: str
    content: str
    modified: bool = False


class PromptValidation(BaseModel):
    ok: bool
    errors: list[str] = Field(default_factory=list)


class StateUpdate(BaseModel):
    keys: list[str]
    values: dict[str, Any]
    full_state: dict[str, Any] | None = None


class ClientMessage(BaseModel):
    type: str  # subscribe, unsubscribe, command
    payload: dict[str, Any] = Field(default_factory=dict)


class ServerMessage(BaseModel):
    type: str  # state_update, agent_event, metrics_update, error, session_update
    timestamp: datetime = Field(default_factory=datetime.utcnow)
    payload: dict[str, Any] = Field(default_factory=dict)


class ImportSessionRequest(BaseModel):
    folder_path: str
    name: str | None = None


class SessionUpdateRequest(BaseModel):
    name: str | None = None


# Copilot proxy models
class CopilotMessageRole(str, Enum):
    SYSTEM = "system"
    USER = "user"
    ASSISTANT = "assistant"


class CopilotToolCallPart(BaseModel):
    """Tool call embedded in an assistant message (OpenAI-style)."""
    call_id: str
    name: str
    arguments: str  # JSON string


class CopilotToolResultPart(BaseModel):
    """Tool execution result embedded in a user message (OpenAI-style)."""
    call_id: str
    content: str
    is_error: bool | None = None


class CopilotMessage(BaseModel):
    """A single message in the conversation."""
    role: CopilotMessageRole
    content: str
    # Preserve tool metadata; otherwise /api/copilot/pending will strip these fields and
    # VS Code can't correlate tool calls/results (causing repeated tool calls).
    tool_calls: list[CopilotToolCallPart] | None = None
    tool_results: list[CopilotToolResultPart] | None = None


class CopilotToolDefinition(BaseModel):
    """Tool/function definition for the model."""
    name: str
    description: str
    parameters: dict[str, Any] = Field(default_factory=dict)


class CopilotToolCall(BaseModel):
    """A tool call made by the model."""
    id: str
    name: str
    arguments: str  # JSON string


class CopilotRequest(BaseModel):
    """LLM request to be processed by Copilot via VS Code extension."""
    request_id: str
    messages: list[CopilotMessage]
    model: str
    tools: list[CopilotToolDefinition] | None = None
    temperature: float | None = None
    max_tokens: int | None = None
    created_at: datetime = Field(default_factory=datetime.utcnow)


class CopilotResponse(BaseModel):
    """Response from Copilot via VS Code extension."""
    request_id: str
    content: str | None = None
    tool_calls: list[CopilotToolCall] | None = None
    error: str | None = None
    finish_reason: str | None = None


class CopilotModelInfo(BaseModel):
    """Information about an available Copilot model."""
    id: str
    name: str
    vendor: str
    family: str
    version: str
    max_input_tokens: int | None = None


class ParetoFrontResponse(BaseModel):
    """Response model for Pareto front data."""

    pareto_front: list[str]
    """Candidate IDs on the Pareto front."""

    primary: str | None
    """Primary selected candidate (highest avg on front)."""

    dominated: list[str]
    """Candidate IDs that are dominated."""
