import logging
from typing import Dict, Any, Optional, Tuple, Union

from sfn_blueprint.utils.llm_handler import SFNAIHandler

from .config import GluonConfig
from .constants import (
    format_autogluon_config_prompt,
    AUTOGLUON_CONFIG_SYSTEM_PROMPT,
    AUTOGLUON_STRUCTURING_SYSTEM_PROMPT,
    RECOMMENDATION_STRUCTURING_SYSTEM_PROMPT,
    CLASSIFICATION_STRUCTURING_PROMPT,
    REGRESSION_STRUCTURING_PROMPT,
    FORECASTING_STRUCTURING_PROMPT,
    ANOMALY_DETECTION_STRUCTURING_PROMPT,
    CLUSTERING_STRUCTURING_PROMPT,
)
from .models import (
    AutoGluonConfig,
    ClassificationUserExperimentConfig,
    RegressionUserExperimentConfig,
    ForecastingUserExperimentConfig,
    AnomalyDetectionUserExperimentConfig,
    RecommendationUserExperimentConfig,
    ClusteringUserExperimentConfig
)


class AutoGluonConfigAgent:
    """
    Two-stage AutoGluon configuration agent:

    Agent 1 → Reasoning / recommendation (human-style)
    Agent 2 → Structuring / strict JSON (Pydantic enforced)
    """

    def __init__(self, config: Optional[GluonConfig] = None):
        self.config = config or GluonConfig()
        self.logger = logging.getLogger(__name__)
        self.ai_handler = SFNAIHandler(logger_name="AutoGluonConfigAgent")

    def _run_reasoning_agent(
        self,
        domain: Dict[str, Any],
        use_case: Dict[str, Any],
        methodology: str,
        dataset_insights: Dict[str, Any],
    ) -> Tuple[str, Dict[str, Any]]:

        system_prompt, user_prompt = format_autogluon_config_prompt(
            domain=domain,
            use_case=use_case,
            methodology=methodology,
            dataset_insights=dataset_insights,
        )

        response, cost = self.ai_handler.route_to(
            llm_provider=self.config.provider,
            model=self.config.model,
            configuration={
                "messages": [
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt},
                ],
                "temperature": self.config.temperature,
                "max_tokens": self.config.max_tokens,
            },
        )
        
        return response, cost


    def _get_structuring_prompt_and_model(self, methodology: str) -> Tuple[str, type]:
        """
        Get the appropriate structuring prompt and model class based on methodology.

        Returns:
            Tuple of (structuring_prompt, model_class)
        """
        methodology_lower = methodology.lower()

        # Classification methodologies
        if any(keyword in methodology_lower for keyword in ['binary_classification', 'multi_classification', 'multiclass_classification', 'timeseries_binary', 'timeseries_multiclass']):
            return CLASSIFICATION_STRUCTURING_PROMPT, ClassificationUserExperimentConfig

        # Regression methodologies
        elif any(keyword in methodology_lower for keyword in ['regression', 'timeseries_regression']):
            return REGRESSION_STRUCTURING_PROMPT, RegressionUserExperimentConfig

        # Forecasting
        elif 'forecasting' in methodology_lower:
            return FORECASTING_STRUCTURING_PROMPT, ForecastingUserExperimentConfig

        # Anomaly detection
        elif 'anomaly' in methodology_lower:
            return ANOMALY_DETECTION_STRUCTURING_PROMPT, AnomalyDetectionUserExperimentConfig

        # Recommendation
        elif 'recommendation' in methodology_lower:
            return RECOMMENDATION_STRUCTURING_SYSTEM_PROMPT, RecommendationUserExperimentConfig

        # Clustering
        elif 'clustering' in methodology_lower:
            return CLUSTERING_STRUCTURING_PROMPT, ClusteringUserExperimentConfig

        # Default to legacy AutoGluon config
        else:
            return AUTOGLUON_STRUCTURING_SYSTEM_PROMPT, AutoGluonConfig

    def _run_structuring_agent(
        self,
        reasoning_output: str,
        methodology: str,
    ) -> Tuple[Union[ClassificationUserExperimentConfig, RegressionUserExperimentConfig, ForecastingUserExperimentConfig,
                      AnomalyDetectionUserExperimentConfig, RecommendationUserExperimentConfig, ClusteringUserExperimentConfig,
                      AutoGluonConfig], Dict[str, Any]]:
        """
        Run the structuring agent to convert reasoning output to structured JSON.
        Automatically selects the correct prompt and model based on methodology.
        """
        structuring_prompt, model_class = self._get_structuring_prompt_and_model(methodology)

        response, cost = self.ai_handler.route_to(
            llm_provider=self.config.provider,
            model=self.config.model,
            configuration={
                "messages": [
                    {
                        "role": "system",
                        "content": structuring_prompt,
                    },
                    {
                        "role": "user",
                        "content": reasoning_output,
                    },
                ],
                "temperature": self.config.temperature,
                "max_tokens": self.config.max_tokens,
            },
        )

        config = model_class.model_validate_json(response)

        return config, cost

    def _run_recommendation_structuring_agent(
        self,
        reasoning_output: str,
    ) -> Tuple[RecommendationUserExperimentConfig, Dict[str, Any]]:
        """Legacy method - kept for backward compatibility"""
        response, cost = self.ai_handler.route_to(
            llm_provider=self.config.provider,
            model=self.config.model,
            configuration={
                "messages": [
                    {
                        "role": "system",
                        "content": RECOMMENDATION_STRUCTURING_SYSTEM_PROMPT,
                    },
                    {
                        "role": "user",
                        "content": reasoning_output,
                    },
                ],
                "temperature": self.config.temperature,
                "max_tokens": self.config.max_tokens,
            },
        )

        config = RecommendationUserExperimentConfig.model_validate_json(response)

        return config, cost    


    def configure_training(
        self,
        domain: Dict[str, Any],
        use_case: Dict[str, Any],
        methodology: str,
        dataset_insights: Dict[str, Any],
    ) -> Tuple[Union[ClassificationUserExperimentConfig, RegressionUserExperimentConfig, ForecastingUserExperimentConfig,
                      AnomalyDetectionUserExperimentConfig, RecommendationUserExperimentConfig, ClusteringUserExperimentConfig,
                      AutoGluonConfig], Dict[str, Any]]:
        """
        Configure training parameters based on domain, use case, methodology, and dataset.

        Returns:
            Tuple of (config, cost_summary) where config is the appropriate methodology-specific
            configuration object with user_experiment_config section.
        """
        reasoning_output, cost_1 = self._run_reasoning_agent(
            domain, use_case, methodology, dataset_insights
        )

        # Use the new unified structuring agent that handles all methodologies
        config, cost_2 = self._run_structuring_agent(reasoning_output, methodology)

        return config, {
            "reasoning_agent": cost_1,
            "structuring_agent": cost_2,
        }

    def execute_task(self, task_data: Dict[str, Any]) -> Dict[str, Any]:
        config, cost = self.configure_training(
            domain=task_data["domain"],
            use_case=task_data["use_case"],
            methodology=task_data["methodology"],
            dataset_insights=task_data["dataset_insights"],
        )

        return {
            "configuration": config.model_dump(),
            "cost_summary": cost,
        }

    def __call__(self, task_data: Dict[str, Any]) -> Dict[str, Any]:
        return self.execute_task(task_data)