"""Configuration for Agent Cost Optimizer."""

from dataclasses import dataclass, field
from typing import Dict, List, Optional, Any
from pathlib import Path
import yaml


@dataclass
class ModelConfig:
    model_id: str
    provider: str
    cost_per_1k_input: float
    cost_per_1k_output: float
    cost_per_1k_reasoning: float = 0.0
    latency_ms_estimate: float = 1000.0
    strength_tier: int = 3  # 1=tiny, 2=cheap, 3=medium, 4=frontier, 5=specialist, 6=verifier
    max_context: int = 128000
    supports_tools: bool = True
    supports_reasoning: bool = False
    cache_discount_rate: float = 0.5


@dataclass
class ToolConfig:
    tool_name: str
    cost_per_call: float = 0.0
    latency_ms_estimate: float = 500.0
    cacheable: bool = False
    requires_verification: bool = False
    max_retries: int = 3


@dataclass
class VerifierConfig:
    verifier_model_id: str
    cost_per_call: float = 0.0
    latency_ms_estimate: float = 1000.0
    confidence_threshold: float = 0.8


@dataclass
class RoutingPolicy:
    name: str
    type: str = "cascade"  # cascade, static, learned, prompt_only
    threshold_confidence: float = 0.7
    max_cascade_depth: int = 3
    enable_verifier_fallback: bool = True
    enable_escalation: bool = True


@dataclass
class ACOConfig:
    project_name: str = "agent-cost-optimizer"
    trace_storage_path: str = "./traces"
    models: Dict[str, ModelConfig] = field(default_factory=dict)
    tools: Dict[str, ToolConfig] = field(default_factory=dict)
    verifiers: Dict[str, VerifierConfig] = field(default_factory=dict)
    routing_policy: RoutingPolicy = field(default_factory=lambda: RoutingPolicy("default"))
    
    # Cost weights
    model_cost_weight: float = 1.0
    tool_cost_weight: float = 1.0
    verifier_cost_weight: float = 1.0
    latency_weight: float = 0.1
    retry_penalty_weight: float = 2.0
    false_done_penalty: float = 10.0
    unsafe_cheap_model_penalty: float = 20.0
    missed_escalation_penalty: float = 15.0
    
    # Module toggles
    enable_telemetry: bool = True
    enable_classifier: bool = True
    enable_router: bool = True
    enable_context_budgeter: bool = True
    enable_cache_layout: bool = True
    enable_tool_gate: bool = True
    enable_verifier_budgeter: bool = True
    enable_retry_optimizer: bool = True
    enable_meta_tool_miner: bool = True
    enable_early_termination: bool = True
    
    # Cache-aware layout
    cache_prefix_stable: List[str] = field(default_factory=lambda: [
        "system_rules", "tool_descriptions", "user_preferences"
    ])
    cache_suffix_dynamic: List[str] = field(default_factory=lambda: [
        "user_message", "retrieved_docs", "recent_trace", "artifacts"
    ])
    
    # Early termination
    doom_max_cost_ratio: float = 3.0  # stop if cost > 3x predicted
    doom_max_retries: int = 3
    doom_no_progress_steps: int = 5
    doom_verifier_disagreement_threshold: int = 2
    
    # Meta-tool mining
    meta_tool_min_frequency: int = 5
    meta_tool_min_success_rate: float = 0.8
    
    @classmethod
    def from_yaml(cls, path: str) -> "ACOConfig":
        with open(path, "r") as f:
            data = yaml.safe_load(f)
        
        models = {k: ModelConfig(**v) for k, v in data.get("models", {}).items()}
        tools = {k: ToolConfig(**v) for k, v in data.get("tools", {}).items()}
        verifiers = {k: VerifierConfig(**v) for k, v in data.get("verifiers", {}).items()}
        routing = RoutingPolicy(**data.get("routing_policy", {}))
        
        config = cls(
            project_name=data.get("project_name", "agent-cost-optimizer"),
            trace_storage_path=data.get("trace_storage_path", "./traces"),
            models=models,
            tools=tools,
            verifiers=verifiers,
            routing_policy=routing,
        )
        
        # Override any direct fields
        for key in ["model_cost_weight", "tool_cost_weight", "verifier_cost_weight",
                    "latency_weight", "retry_penalty_weight", "false_done_penalty",
                    "unsafe_cheap_model_penalty", "missed_escalation_penalty",
                    "enable_telemetry", "enable_classifier", "enable_router",
                    "enable_context_budgeter", "enable_cache_layout", "enable_tool_gate",
                    "enable_verifier_budgeter", "enable_retry_optimizer",
                    "enable_meta_tool_miner", "enable_early_termination",
                    "doom_max_cost_ratio", "doom_max_retries", "doom_no_progress_steps",
                    "doom_verifier_disagreement_threshold",
                    "meta_tool_min_frequency", "meta_tool_min_success_rate"]:
            if key in data:
                setattr(config, key, data[key])
        
        return config
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "project_name": self.project_name,
            "trace_storage_path": self.trace_storage_path,
            "models": {k: vars(v) for k, v in self.models.items()},
            "tools": {k: vars(v) for k, v in self.tools.items()},
            "verifiers": {k: vars(v) for k, v in self.verifiers.items()},
            "routing_policy": vars(self.routing_policy),
        }