from typing import Dict, Any, Optional, List
from pydantic import BaseModel, Field, model_validator
from eveai_chat_workers.retrievers.retriever_typing import RetrieverArguments
from common.extensions import cache_manager


class SpecialistArguments(BaseModel):
    """
    Dynamic arguments for specialists, allowing arbitrary fields but validating required ones
    based on SPECIALIST_TYPES configuration.
    """
    type: str = Field(..., description="Type of specialist (e.g. STANDARD_RAG)")
    type_version: str = Field(..., description="Type version (e.g. 1.0)")
    retriever_arguments: Dict[str, Any] = Field(
        default_factory=dict,
        description="Arguments for each retriever, keyed by retriever ID"
    )

    # Allow any additional fields
    model_config = {
        "extra": "allow"
    }

    # Structural optional fields available for all specialists
    question: Optional[str] = Field(
        None,
        description="Optional question directed to the specialist"
    )
    form_values: Optional[Dict[str, Any]] = Field(
        None,
        description="Optional form values filled by the user, keyed by field name"
    )

    @model_validator(mode='after')
    def validate_required_arguments(self) -> 'SpecialistArguments':
        """Validate that all required arguments for this specialist type are present"""
        specialist_config = cache_manager.specialists_config_cache.get_config(self.type, self.type_version)
        if not specialist_config:
            raise ValueError(f"Unknown specialist type: {self.type}")

        # Check required arguments from configuration
        for arg_name, arg_config in specialist_config['arguments'].items():
            if arg_config.get('required', False):
                if not hasattr(self, arg_name):
                    raise ValueError(f"Missing required argument '{arg_name}' for {self.type}")

                # Type validation
                value = getattr(self, arg_name)
                expected_type = arg_config['type']
                if expected_type == 'str' and not isinstance(value, str):
                    raise ValueError(f"Argument '{arg_name}' must be a string")
                elif expected_type == 'int' and not isinstance(value, int):
                    raise ValueError(f"Argument '{arg_name}' must be an integer")

        return self

    @classmethod
    def create(cls, type_name: str, type_version: str, specialist_args: Dict[str, Any],
               retriever_args: Dict[str, Dict[str, Any]]) -> 'SpecialistArguments':
        """
        Factory method to create SpecialistArguments with validated retriever arguments

        Args:
            type_name: The specialist type (e.g., 'STANDARD_RAG')
            specialist_args: Arguments specific to the specialist
            retriever_args: Dictionary of retriever arguments keyed by retriever ID

        Returns:
            Validated SpecialistArguments instance
        """
        # Convert raw retriever arguments to RetrieverArguments instances
        validated_retriever_args = {}
        for retriever_id, args in retriever_args.items():
            # Ensure type is included in retriever arguments
            if 'type' not in args:
                raise ValueError(f"Retriever arguments for {retriever_id} must include 'type'")
            if 'type_version' not in args:
                raise ValueError(f"Retriever arguments for {retriever_id} must include 'type_version'")

            validated_retriever_args[retriever_id] = RetrieverArguments(**args)

        # Combine everything into the specialist arguments
        return cls(
            type=type_name,
            type_version=type_version,
            **specialist_args,
            retriever_arguments=validated_retriever_args
        )


class SpecialistResult(BaseModel):
    """
    Dynamic results from specialists, validating required fields based on
    SPECIALIST_TYPES configuration.
    """
    type: str = Field(..., description="Type of specialist (e.g. STANDARD_RAG)")
    type_version: str = Field(..., description="Type version (e.g. 1.0)")

    # Allow any additional fields
    model_config = {
        "extra": "allow"
    }

    # Structural optional fields available for all specialists
    answer: Optional[str] = Field(None, description="Optional textual answer from the specialist")
    form_request: Optional[Dict[str, Any]] = Field(None, description="Optional form definition to request user input")
    phase: Optional[str] = Field(None, description="Phase of the specialist's workflow")
    citations: Optional[List[Dict[str, Any]]] = Field(None, description="Citations for the specialist's answer")

    @model_validator(mode='after')
    def validate_required_results(self) -> 'SpecialistResult':
        """Validate that all required result fields for this specialist type are present"""
        specialist_config = cache_manager.specialists_config_cache.get_config(self.type, self.type_version)
        if not specialist_config:
            raise ValueError(f"Unknown specialist type: {self.type}, {self.type_version}")

        # Check required results from configuration
        required_results = specialist_config.get('results', {})
        for result_name, result_config in required_results.items():
            if result_config.get('required', False):
                if not hasattr(self, result_name):
                    raise ValueError(f"Missing required result '{result_name}' for {self.type}")

                # Type validation
                value = getattr(self, result_name)
                expected_type = result_config['type']

                # Validate based on type annotation
                if expected_type == 'str' and not isinstance(value, str):
                    raise ValueError(f"Result '{result_name}' must be a string")
                elif expected_type == 'bool' and not isinstance(value, bool):
                    raise ValueError(f"Result '{result_name}' must be a boolean")
                elif expected_type == 'List[str]' and not (
                        isinstance(value, list) and all(isinstance(x, str) for x in value)):
                    raise ValueError(f"Result '{result_name}' must be a list of strings")
                # Add other type validations as needed

        return self

    @classmethod
    def create_for_type(cls, specialist_type: str, specialist_type_version: str, **results) -> 'SpecialistResult':
        """
        Factory method to create a type-specific result

        Args:
            specialist_type: The type of specialist (e.g., 'STANDARD_RAG')
            specialist_type_version: The type of specialist (e.g., '1.0')
            **results: The result values to include

        Returns:
            Validated SpecialistResult instance

        Example:
            For STANDARD_RAG:
            result = SpecialistResult.create_for_type(
                'STANDARD_RAG',
                '1.0',
                answer="The answer text",
                citations=["doc1", "doc2"],
                insufficient_info=False
            )
        """
        # Add the type to the results
        results['type'] = specialist_type
        results['type_version'] = specialist_type_version

        # Create and validate the result
        return cls(**results)
