import os
from typing import Dict, List, Any, Optional
import logging
from datetime import datetime
from environs import Env

from llmops.preprocessor import Preprocessor
from llmops.subtask_orchestrator import SubTaskOrchestrator, TaskConfig
from llmops.storage.transition_store import TransitionStore
from llmops.config import EnvironmentConfig, AgentConfig
from llmops.business.business_expert import BusinessExpert


class TaskOrchestrator:
    """whole task orchestrator"""

    def __init__(
        self,
        output_dir: str,
        base_url: str,
        max_steps: int,
        agent_type: str,
        epsilon: float,
        llm_model: str,
        llm_temperature: float,
        transition_store: TransitionStore,
        business_expert: BusinessExpert,
    ):
        self.output_dir = output_dir
        self.base_url = base_url
        self.max_steps = max_steps
        self.agent_type = agent_type
        self.epsilon = epsilon
        self.llm_model = llm_model
        self.llm_temperature = llm_temperature
        self.transition_store = transition_store
        self.business_expert = business_expert
        self.preprocessor = Preprocessor(self.business_expert, llm_model)

    def execute_scenarios(
        self,
        scenarios: Dict[str, List[str]],
        sut_token: str | Dict[str, str],
        auth_type: Optional[str] = None,
    ):
        """execute all scenarios

        Args:
            scenarios: scenario dictionary, key is scenario name, value is scenario steps
            sut_token: SUT token
        """
        for scenario_name, steps in scenarios.items():
            logging.info(f"\nExecuting scenario: {scenario_name}")

            # generate abstract state sequence for current scenario
            abstract_ops = self.preprocessor.process(steps)
            logging.info(
                f"Generated {len(abstract_ops)} abstract operations for scenario {scenario_name}"
            )

            # create task id
            generation_id = (
                f"{datetime.now().strftime('%Y%m%d_%H%M%S')}_{scenario_name}"
            )

            # update environment configuration
            env_config = EnvironmentConfig(
                base_url=self.base_url,
                api_token=sut_token,
                auth_type=auth_type if auth_type else "bearer",
                generation_id=generation_id,
            )

            # create task config
            task_config = TaskConfig(
                env_config=env_config,
                max_steps=self.max_steps,
            )

            # create agent config
            agent_config = AgentConfig(
                agent_type=self.agent_type,
                epsilon=self.epsilon,
                llm_model=self.llm_model,
                llm_temperature=self.llm_temperature,
            )

            # create sub task orchestrator
            orchestrator = SubTaskOrchestrator(
                config=task_config,
                agent_config=agent_config,
                business_expert=self.business_expert,
                transition_store=self.transition_store,
            )

            # execute scenario
            report = orchestrator.execute(abstract_ops)

            # save report
            report_file = os.path.join(self.output_dir, f"{generation_id}_report.json")
            with open(report_file, "w") as f:
                # use pydantic's json method, it handles datetime etc.
                json_str = report.model_dump_json(indent=2)
                f.write(json_str)

            logging.info(f"Scenario {scenario_name}:")
            logging.info(f"- Success: {'Yes' if report.success else 'No'}")
            logging.info(f"- Total steps: {report.steps_taken}")
            logging.info(f"- Total episodes: {report.episodes}")
            if report.error:
                logging.info(f"- Error: {report.error}")
            logging.info(
                f"- Execution time: {(report.end_time - report.start_time).total_seconds():.2f}s"
            )