"""Configuration for Deep Research workflow integration."""

import os
from typing import Optional
from pydantic import BaseModel, Field
from langchain_core.runnables import RunnableConfig


class DeepResearchConfig(BaseModel):
    """Configuration for Deep Research agent workflow."""
    
    # General Configuration
    max_structured_output_retries: int = Field(
        default=3,
        description="Maximum number of retries for structured output calls from models"
    )
    allow_clarification: bool = Field(
        default=False,  # Disabled by default for evaluation tasks
        description="Whether to allow the researcher to ask the user clarifying questions"
    )
    max_concurrent_research_units: int = Field(
        default=3,  # Conservative default for evaluation
        description="Maximum number of research units to run concurrently"
    )
    
    # Research Configuration
    max_researcher_iterations: int = Field(
        default=4,  # Conservative for math problems
        description="Maximum number of research iterations for the Research Supervisor"
    )
    max_react_tool_calls: int = Field(
        default=8,  # Conservative for math problems
        description="Maximum number of tool calling iterations to make in a single researcher step"
    )
    
    # Model Configuration
    research_model: str = Field(
        default="gpt-4o",
        description="Model for conducting research"
    )
    research_model_max_tokens: int = Field(
        default=4096,
        description="Maximum output tokens for research model"
    )
    research_model_temperature: float = Field(
        default=0.7,
        description="Temperature for research model"
    )
    compression_model: str = Field(
        default="gpt-4o",
        description="Model for compressing research findings from sub-agents"
    )
    compression_model_max_tokens: int = Field(
        default=8192,  # Increased to account for GPT-5 reasoning tokens
        description="Maximum output tokens for compression model"
    )
    compression_model_temperature: float = Field(
        default=0.7,
        description="Temperature for compression model"
    )
    final_report_model: str = Field(
        default="gpt-4o",
        description="Model for writing the final report from all research findings"
    )
    final_report_model_max_tokens: int = Field(
        default=4096,
        description="Maximum output tokens for final report model"
    )
    final_report_model_temperature: float = Field(
        default=0.7,
        description="Temperature for final report model"
    )
    
    # Model Provider Configuration
    research_model_provider: Optional[str] = Field(
        default=None,
        description="Provider for research model (e.g., 'openai', 'anthropic')"
    )
    compression_model_provider: Optional[str] = Field(
        default=None,
        description="Provider for compression model (e.g., 'openai', 'anthropic')"
    )
    final_report_model_provider: Optional[str] = Field(
        default=None,
        description="Provider for final report model (e.g., 'openai', 'anthropic')"
    )
    
    # Search Integration (uses existing search engine)
    enable_web_search: bool = Field(
        default=False,
        description="Enable web search tools (disabled by default for ChromaDB-only evaluation)"
    )
    
    @classmethod
    def from_model_config(cls, model_config: dict, dr_config: Optional[dict] = None) -> "DeepResearchConfig":
        """Create DeepResearchConfig from model configuration."""
        # Extract relevant fields from model config
        config_data = {
            "research_model": model_config.get("name", "gpt-4o"),
            "compression_model": model_config.get("name", "gpt-4o"), 
            "final_report_model": model_config.get("name", "gpt-4o"),
            "research_model_max_tokens": model_config.get("max_tokens", 4096),
            "compression_model_max_tokens": model_config.get("max_tokens", 4096),
            "final_report_model_max_tokens": model_config.get("max_tokens", 4096),
            "research_model_temperature": model_config.get("temperature", 0.7),
            "compression_model_temperature": model_config.get("temperature", 0.7),
            "final_report_model_temperature": model_config.get("temperature", 0.7),
            "research_model_provider": model_config.get("provider"),
            "compression_model_provider": model_config.get("provider"),
            "final_report_model_provider": model_config.get("provider"),
        }
        # Add any deep research specific overrides from model_config
        if "deep_research_config" in model_config:
            config_data.update(model_config["deep_research_config"])
        if dr_config:
            config_data.update(dr_config)
        
        return cls(**config_data)
    
    @classmethod
    def from_runnable_config(cls, config: Optional[RunnableConfig] = None) -> "DeepResearchConfig":
        """Create a DeepResearchConfig instance from a RunnableConfig."""
        configurable = config.get("configurable", {}) if config else {}
        field_names = list(cls.model_fields.keys())
        values = {
            field_name: os.environ.get(field_name.upper(), configurable.get(field_name))
            for field_name in field_names
        }
        return cls(**{k: v for k, v in values.items() if v is not None})

    class Config:
        """Pydantic configuration."""
        arbitrary_types_allowed = True
