from token import OP
from typing import List, Dict, Any, Optional
from pydantic import BaseModel
from datetime import datetime
import uuid

from llmops.environment.environment import Environment
from llmops.config import EnvironmentConfig, AgentConfig
from llmops.preprocessor import AbstractOperation
from llmops.storage.transition_store import TransitionStore, Transition
from llmops.agents.agent import LLMAgent
from llmops.business.business_expert import BusinessExpert


class TaskConfig(BaseModel):
    """task config"""

    env_config: EnvironmentConfig
    max_steps: int = 20  # max steps per subtask


class ExecutionReport(BaseModel):
    """Execution report for a task"""

    # Task Identification & Status
    task_id: str
    success: bool
    error: Optional[str] = None

    # Timing
    start_time: datetime
    end_time: datetime

    # Statistics
    steps_taken: int
    episodes: int
    final_state: Dict[str, Any]

    # Transitions
    transitions: List[Dict[str, Any]]  # 每个 transition 的 to_dict() 结果

    class Config:
        """Pydantic config"""

        json_encoders = {datetime: lambda v: v.isoformat()}


class SubTaskOrchestrator:
    """sub task orchestrator

    responsible for :
    1. managing task lifecycle
    2. implementing the two-layer traversal and state achievement process
    3. handling resource cleanup
    4. generating execution reports
    """

    def __init__(
        self,
        config: TaskConfig,
        agent_config: AgentConfig,
        business_expert: BusinessExpert,
        transition_store: Optional[TransitionStore] = None,
    ):
        """initialize sub task orchestrator

        Args:
            config: task config
            agent_config: agent config
            transition_store: optional, state transition store. if not provided, a new instance will be created
            business_expert: business expert instance
        """
        self.config = config
        self.agent_config = agent_config  # save agent config for later use
        self.task_id = str(uuid.uuid4())
        self.transition_store = transition_store or TransitionStore()
        self.business_expert = business_expert

        # initialize environment
        env_config = config.env_config
        env_config.generation_id = self.task_id
        self.env = Environment(
            config=env_config,
            business_expert=self.business_expert,
            model=self.agent_config.llm_model,
        )

        # create agent
        self.agent = LLMAgent(business_expert=self.business_expert, model=self.agent_config.llm_model, transition_store=self.transition_store)

        # execution statistics
        self.total_steps = 0
        self.episodes = 0

    def execute(self, abstract_ops: List[AbstractOperation]) -> ExecutionReport:
        """Run an e2e scenario
        Args:
            abstract_ops: a list of AbstractOperation objects from ONE e2e scenario

        Returns:
            ExecutionReport: steps
        """
        start_time = datetime.now()
        # current_state = self.business_expert.list_all_resources()
        current_state = self.env.resource_manager.list_all_resources()  # initialize current state
        self.env.current_state = current_state
        success = True  # initialize as True

        try:
            # Layer 1: State Traversal Loop
            for target_operation in abstract_ops:
                self.env.set_target(target_operation)

                # Layer 2: State Achievement Loop
                try:
                    current_state = self._achieve_operation(target_operation)
                    if not current_state:
                        success = False  # if any state achievement fails, the whole process fails
                        raise Exception(f"Failed to achieve state: {target_operation}")
                except Exception as e:
                    import traceback

                    # generate error report
                    error_traceback = traceback.format_exc()
                    print("错误详情:")
                    print(error_traceback)

                    # Store error transition if we have current state
                    if hasattr(self, "env") and hasattr(self.env, "current_state"):
                        self.transition_store.store(
                            Transition(
                                generation_id=self.task_id,
                                operation_id=self.agent.operation_selector.get_trajectory()[-1]["output"]["operation_id"],
                                state_before=self.env.current_state,
                                state_after=self.env.current_state,  # 错误发生时状态未改变
                                request={},  # 错误发生时可能没有 request
                                response={},  # 错误发生时可能没有 response
                                status_code=-1,  # 使用特殊状态码表示错误
                                error=str(e),
                                error_traceback=error_traceback,
                                timestamp=datetime.now(),
                            )
                        )

                    end_time = datetime.now()
                    return ExecutionReport(
                        task_id=self.task_id,
                        success=False,
                        error=str(e),
                        start_time=start_time,
                        end_time=end_time,
                        steps_taken=self.total_steps,
                        episodes=self.episodes,
                        final_state=current_state if current_state else {},
                        transitions=[
                            t.to_dict()
                            for t in self.transition_store.get_transitions_by_generation(
                                self.task_id
                            )
                        ],
                    )

            # generate success report
            end_time = datetime.now()
            return ExecutionReport(
                task_id=self.task_id,
                success=success,
                start_time=start_time,
                end_time=end_time,
                steps_taken=self.total_steps,
                episodes=self.episodes,
                final_state=current_state,
                transitions=[
                    t.to_dict()
                    for t in self.transition_store.get_transitions_by_generation(
                        self.task_id
                    )
                ],
            )

        except Exception as e:
            import traceback

            # generate error report
            print("错误详情:")
            print(traceback.format_exc())
            end_time = datetime.now()
            return ExecutionReport(
                task_id=self.task_id,
                success=False,
                error=str(e),
                start_time=start_time,
                end_time=end_time,
                steps_taken=self.total_steps,
                episodes=self.episodes,
                final_state=current_state if current_state else {},
                transitions=[
                    t.to_dict()
                    for t in self.transition_store.get_transitions_by_generation(
                        self.task_id
                    )
                ],
            )
        finally:
            # cleanup resources
            # self.business_expert.cleanup()
            self.env.resource_manager.cleanup()

    def _achieve_operation(
        self, target_operation: AbstractOperation
    ) -> Optional[Dict[str, Any]]:
        """achieve target state

        use agent to select actions until reaching the target state or exceeding the maximum attempts

        Args:
            target_operation: target abstract state

        Returns:
            Optional[Dict[str, Any]]: return final state if successful, otherwise return None
        """

        self.episodes += 1
        steps = 0

        # get initial state
        state = self.env.current_state
        done = False
        info = {
            "execution_analyzer": None,
            "attempts": 0,
            "target_matcher": None,
            "response": None,
        }

        while not done and steps <= self.config.max_steps:
            # agent select action
            action = self.agent.get_action(state, target_operation, info)

            try:
                # execute action
                next_state, reward, done, truncated, info = self.env.step(action)

                # Store successful transition
                self.transition_store.store(
                    Transition(
                        generation_id=self.task_id,
                        operation_id=action.operation_id,
                        state_before=state,
                        state_after=next_state,
                        request=self.env.executor.reqeust_data,
                        response=info["response"],
                        status_code=info["response"]["status_code"],
                        operation_selection=self.agent.operation_selector.get_trajectory()[
                            -1
                        ]
                        if hasattr(self.agent, "operation_selector")
                        else None,
                        request_generation=self.agent.request_generator.get_trajectory()[
                            -1
                        ]
                        if hasattr(self.agent, "request_generator")
                        else None,
                        execution_analysis=self.env.execution_analyzer.get_trajectory()[
                            -1
                        ]
                        if hasattr(self.env, "execution_analyzer")
                        else info["execution_analyzer"],
                        target_matching=self.env.target_matcher.get_trajectory()[-1]
                        if hasattr(self.env, "target_matcher")
                        else info["target_match"],
                        resource_manager=self.env.resource_manager.get_trajectory()[-1]
                        if hasattr(self.env, "resource_manager")
                        else None,
                        timestamp=datetime.now(),
                    )
                )
            except Exception as e:
                import traceback

                error_traceback = traceback.format_exc()

                # Store error transition
                self.transition_store.store(
                    Transition(
                        generation_id=self.task_id,
                        operation_id=action.operation_id,
                        state_before=state,
                        state_after=state,  # 错误发生时状态未改变
                        request=self.env.executor.reqeust_data,
                        response={},  # 错误发生时可能没有 response
                        status_code=-1,  # 使用特殊状态码表示错误
                        operation_selection=self.agent.operation_selector.get_trajectory()[
                            -1
                        ]
                        if hasattr(self.agent, "operation_selector")
                        else None,
                        request_generation=self.agent.request_generator.get_trajectory()[
                            -1
                        ]
                        if hasattr(self.agent, "request_generator")
                        else None,
                        execution_analysis=None,  # 错误发生时可能没有执行分析
                        target_matching=None,  # 错误发生时可能没有目标匹配
                        resource_manager=None,  # 错误发生时可能没有资源管理
                        timestamp=datetime.now(),
                        error=str(e),
                        error_traceback=error_traceback,
                    )
                )
                raise  # 重新抛出异常

            # update statistics
            steps += 1
            self.total_steps += 1

            # only update if the agent supports it
            if hasattr(self.agent, "update"):
                self.agent.update(state, action, reward, next_state, done)

            # update state
            state = next_state
            self.env.current_state = state

            done = info["target_match"]["judgement"]
            if done:
                return state

        if not done:
            raise RuntimeError(
                f"Failed to achieve target state within {self.config.max_steps} steps"
            )
        return None