from typing import Dict, Any, Optional, Tuple
from datetime import datetime
from copy import deepcopy

from llmops.config import EnvironmentConfig
from llmops.base import (
    Step,
    ExecutionAnalysisType,
    AbstractOperation,
    StepType,
)
from llmops.environment.executor import CompoundExecutor
from llmops.environment.execution_analyzer import ExecutionAnalyzer
from llmops.environment.target_matcher import TargetMatcher
from llmops.environment.resource_manager import ResourceManager
from llmops.utils.json_utils import format_json_string_by_ai
from llmops.business.business_expert import BusinessExpert


class Environment:
    def __init__(
        self,
        config: EnvironmentConfig,
        business_expert: Optional[BusinessExpert] = None,
        model: str = "gpt-o3-mini",
    ):
        self.config = config
        self.business_expert = business_expert
        self.trajectory = []
        self.executor = CompoundExecutor(
            base_url=config.base_url,
            api_token=config.api_token,
            auth_type=config.auth_type,
        )
        self.execution_analyzer = ExecutionAnalyzer(business_expert, model)
        self.target_matcher = TargetMatcher(business_expert, model)
        self.resource_manager = ResourceManager(business_expert, model)
        self.current_state: Dict[str, Any] = {}
        self.target_operation: Optional[AbstractOperation] = None
        self.generation_id: str = config.generation_id  # meta: generation_id

    def reset(self) -> Dict[str, Any]:
        """Local environment does not support reset"""
        # self.current_state = self.business_expert.list_all_resources()
        self.current_state = self.resource_manager.list_all_resources()

    def set_target(self, target_operation: AbstractOperation):
        """set target operation"""
        self.target_operation = target_operation

    def step(
        self, action: Step
    ) -> Tuple[Dict[str, Any], float, bool, bool, Dict[str, dict]]:
        point = {
            "action": action,
            "prev_state": None,
            "current_state": None,
            "timestamp": datetime.now().strftime("%Y%m%d_%H%M%S"),
            "request": None,
            "response": None,
            "execution_analyzer_result": None,
            "match_result": None,
        }
        self.trajectory.append(point)

        prev_state = deepcopy(self.current_state)
        self.trajectory[-1].update(prev_state=prev_state)

        response = self.executor.execute_with_retry(step=action)

        self.trajectory[-1].update(request=self.executor.requests[-1])
        self.trajectory[-1].update(response=response.model_dump())
        self.last_response = response
        # self.current_state = self.business_expert.list_all_resources()
        self.current_state = self.resource_manager.update_resources(action.request.model_dump(), response.model_dump())
        self.trajectory[-1].update(current_state=self.current_state)
        reward = 0
        # evaluate
        judgement = self.execution_analyzer.evaluate(
            request={"request": action.request.model_dump(), "expect": action.expect},
            response=response.model_dump(),
            prev_state=prev_state,
            current_state=self.current_state,
            target_operation=self.target_operation.model_dump(),
            operation_id=action.operation_id,
        )
        last_judgement = judgement
        self.trajectory[-1].update(execution_analyzer_result=judgement.model_dump())
        target_match_result = self._is_done()

        if last_judgement.result_type == ExecutionAnalysisType.SUCCESS:
            # target_match_result = self._is_done()
            self.trajectory[-1].update(match_result=target_match_result.model_dump())
            terminated = target_match_result.judgement
            truncated = False  # self._is_truncated()
            info = {
                "execution_analyzer": judgement.model_dump(),
                "target_match": target_match_result.model_dump(),
                # "attempts": 0,
                "response": response.model_dump(),
                # "error": None
            }
            if terminated:
                action.step_type = StepType.DEFINED
            return self.current_state, reward, terminated, truncated, info
        elif (
            last_judgement.result_type == ExecutionAnalysisType.OPERATION_ERROR
            or last_judgement.result_type == ExecutionAnalysisType.REQUEST_ERROR
        ):
            terminated = False
            truncated = False
            self.trajectory[-1].update(
                match_result={
                    "judgement": False,
                    "reason": "The last request failed, so the target is not reached.",
                }
            )
            info = {
                "execution_analyzer": judgement.model_dump(),
                "target_match": {
                    "judgement": False,
                    "reason": "The last request failed, so the target is not reached.",
                },
                # "attempts": 0,
                "response": response.model_dump(),
                # "error": None
            }
            return self.current_state, reward, terminated, truncated, info
        else:
            raise ValueError(f"Unknown result type: {judgement.result_type}")

    def _is_done(self) -> bool:
        """Check if the target operation is reached."""
        return self.target_matcher.matches(
            self.trajectory[-1]["action"], self.last_response, self.target_operation
        )