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

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,
)
from .models import AutoGluonConfig


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": AUTOGLUON_CONFIG_SYSTEM_PROMPT},
                    {"role": "user", "content": user_prompt},
                ],
                "temperature": self.config.temperature,
                "max_tokens": self.config.max_tokens,
            },
        )

        return response, cost


    def _run_structuring_agent(
        self,
        reasoning_output: str,
    ) -> Tuple[AutoGluonConfig, Dict[str, Any]]:

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

        config = AutoGluonConfig.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[AutoGluonConfig, Dict[str, Any]]:

        reasoning_output, cost_1 = self._run_reasoning_agent(
            domain, use_case, methodology, dataset_insights
        )

        config, cost_2 = self._run_structuring_agent(reasoning_output)

        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)