import json
import threading
from typing import Any, Optional

from openenv.core.env_server.interfaces import Environment
from openenv.core.env_server.types import State

from cloudenv.models import CloudEnvAction, CloudEnvObservation, CloudEnvState
from cloudenv.server.exceptions import CloudEnvBadRequest, CloudEnvNotFound
from cloudenv.server.managers.pipeline_manager import PipelineManager
from cloudenv.server.managers.reward_manager import RewardManager
from cloudenv.server.models.pipeline import Policy
from cloudenv.server.providers.aws.adapter import AWSAdapter
from cloudenv.server.providers.aws.iam_enforcer import AccessDenied, IAMEnforcer
from cloudenv.server.utils.logging import logger


class CloudEnvironment(Environment):
    SUPPORTS_CONCURRENT_SESSIONS = False

    def __init__(
        self,
        aws: AWSAdapter,
        pipeline_manager: PipelineManager,
    ) -> None:
        super().__init__()
        self._aws_lock = threading.RLock()

        self.aws = aws
        self.pipeline_manager = pipeline_manager
        self.rewarder = RewardManager(self.pipeline_manager)

        self._state: Optional[CloudEnvState] = None

    # --------------------------------------------------
    # SYNC (ONLY AFTER SUCCESS)
    # --------------------------------------------------
    def _sync_iam_to_pipeline(self, action: CloudEnvAction) -> None:
        pipeline = self.pipeline_manager.get_pipeline(action.task_id)
        if not pipeline:
            return

        payload = action.payload or {}
        operation = action.operation

        def parse_policy(raw_policy: Any) -> Optional[dict]:
            if raw_policy is None:
                return None
            if isinstance(raw_policy, str):
                try:
                    return json.loads(raw_policy)
                except Exception:
                    logger.warning(
                        f"[SYNC] Failed to parse policy: {raw_policy}"
                    )
                    return None
            if isinstance(raw_policy, dict):
                return raw_policy
            return None

        def update_nodes_for_source(role_name: str, policy: dict) -> None:
            actions = []
            for stmt in policy.get("Statement", []):
                acts = stmt.get("Action", [])
                if isinstance(acts, str):
                    acts = [acts]
                actions.extend(acts)

            for src_key, nodes in pipeline.broken_graph.items():
                _, src_instance = (
                    src_key.split(":", 1)
                    if ":" in src_key
                    else (src_key, src_key)
                )

                if src_instance != role_name:
                    continue

                for node in nodes:
                    # Only apply if action matches target service
                    if any(a.startswith(node.service + ":") for a in actions):
                        node.policies = [Policy(**policy)]

        def update_nodes_for_resource(
            service_id: str, instance_id: str, policy: dict
        ) -> None:
            for nodes in pipeline.broken_graph.values():
                for node in nodes:
                    if (
                        node.service == service_id
                        and node.instance == instance_id
                    ):
                        node.policies = [Policy(**policy)]

        logger.debug(f"[SYNC] Processing {action.service}:{operation}")

        if action.service == "iam" and operation == "PutRolePolicy":
            role_name = payload.get("RoleName")
            policy_doc = parse_policy(payload.get("PolicyDocument"))
            if role_name and policy_doc:
                update_nodes_for_source(role_name, policy_doc)

        elif action.service == "s3" and operation == "PutBucketPolicy":
            bucket = payload.get("Bucket")
            policy_doc = parse_policy(payload.get("Policy"))
            if bucket and policy_doc:
                update_nodes_for_resource("s3", bucket, policy_doc)

        elif action.service == "sqs" and operation == "SetQueueAttributes":
            queue_url = payload.get("QueueUrl")
            attrs = payload.get("Attributes", {}) or {}
            policy_doc = parse_policy(attrs.get("Policy"))
            if queue_url and policy_doc:
                update_nodes_for_resource("sqs", queue_url, policy_doc)

        elif action.service == "sns" and operation == "SetTopicAttributes":
            topic_arn = payload.get("TopicArn")
            attr_name = payload.get("AttributeName")
            policy_doc = parse_policy(payload.get("AttributeValue"))
            if topic_arn and str(attr_name) == "Policy" and policy_doc:
                update_nodes_for_resource("sns", topic_arn, policy_doc)

        elif action.service == "dynamodb" and operation == "PutResourcePolicy":
            table = payload.get("ResourceArn") or payload.get("TableName")
            policy_doc = parse_policy(payload.get("Policy"))
            if table and policy_doc:
                update_nodes_for_resource("dynamodb", table, policy_doc)

        logger.debug(f"[SYNC] Completed for {action.service}:{operation}")

    # --------------------------------------------------
    # RESET
    # --------------------------------------------------
    def reset(
        self,
        seed: Optional[int] = None,
        episode_id: Optional[str] = None,
        **kwargs: Any,
    ) -> CloudEnvObservation:

        # if not episode_id:
        #     raise CloudEnvBadRequest("reset requires episode_id")

        task_id = kwargs.get("task_id", None)
        logger.info(f"[RESET] task={task_id}, episode={episode_id}")

        self.aws.reset()
        result = self.aws.get_last_operation_result()

        return CloudEnvObservation(
            episode_id=episode_id,
            task_id=task_id,
            step_count=0,
            operation_result=result,
            reward=0.0,
            done=False,
            current_pipeline_state={},
        )

    # --------------------------------------------------
    # STEP
    # --------------------------------------------------
    def step(
        self,
        action: CloudEnvAction,
        timeout_s: Optional[float] = None,
        **kwargs: Any,
    ) -> CloudEnvObservation:

        logger.info(
            f"[STEP] task={action.task_id}, episode={action.episode_id}"
        )

        if action.operation == "InitializeEnvironment":
            print("REACHED")
            self._state = CloudEnvState(
                task_id=action.task_id,
                episode_id=action.episode_id,
                step_count=0,
                reward_history=[],
            )

            return CloudEnvObservation(
                episode_id=self._state.episode_id,
                task_id=self._state.task_id,
                step_count=self._state.step_count,
                operation_result={},
                reward=0.0,
                done=False,
                current_pipeline_state={},
            )

        if not action.episode_id or not action.task_id:
            raise CloudEnvBadRequest("action requires episode_id and task_id")

        if action.task_id != self._state.task_id:
            raise CloudEnvBadRequest("task mismatch")

        pipeline = self.pipeline_manager.get_pipeline(action.task_id)
        if pipeline is None:
            raise CloudEnvNotFound(
                f"No pipeline found for task_id={action.task_id}"
            )

        logger.info(
            f"[STEP] task={action.task_id}, episode={action.episode_id}, "
            f"service={action.service}, operation={action.operation}"
        )

        operation = action.operation

        # --------------------------------------------------
        # CONFIG CRUD
        # --------------------------------------------------
        if operation in ["AddConfig", "UpdateConfig", "DeleteConfig"]:
            logger.debug(
                f"[CONFIG] operation={operation}, payload={action.payload}"
            )

            if operation == "AddConfig":
                pipeline.add_node_config(
                    action.service,
                    action.instance_id,
                    action.payload.get("config", {}),
                )
            elif operation == "UpdateConfig":
                pipeline.update_node_config(
                    action.service,
                    action.instance_id,
                    action.payload.get("config", {}),
                )
            elif operation == "DeleteConfig":
                pipeline.delete_node_config(
                    action.service,
                    action.instance_id,
                    action.payload.get("keys", []),
                )

            reward, done = self.rewarder.compute_reward(
                pipeline=pipeline, action=action, state=self._state
            )

            self._state.step_count += 1

            logger.info(
                f"[STEP RESULT] step={self._state.step_count}, reward={reward}, done={done}"
            )

            return CloudEnvObservation(
                episode_id=self._state.episode_id,
                task_id=self._state.task_id,
                step_count=self._state.step_count,
                operation_result={"operation": operation, "status": "SUCCESS"},
                reward=reward,
                done=done,
                current_pipeline_state=pipeline.broken_graph,
            )

        # --------------------------------------------------
        # AWS EXECUTION
        # --------------------------------------------------
        with self._aws_lock:
            success = True

            try:
                logger.debug(f"[AWS EXECUTE] payload={action.payload}")

                self.aws.execute(
                    episode_id=action.episode_id,
                    task_id=action.task_id,
                    service=action.service,
                    action=action.operation,
                    payload=action.payload,
                    principal_arn=action.principal_arn,
                )

                logger.info(
                    f"[AWS SUCCESS] {action.service}:{action.operation}"
                )

            except AccessDenied as e:
                success = False
                logger.warning(
                    f"[ACCESS DENIED] {action.service}:{action.operation} error={str(e)}"
                )

                self.aws.record_operation_failure(
                    episode_id=action.episode_id,
                    task_id=action.task_id,
                    service=action.service,
                    operation=action.operation,
                    error=str(e),
                    error_type=type(e).__name__,
                    principal=action.principal_arn
                    or IAMEnforcer.DEFAULT_ROOT_PRINCIPAL,
                )

            except Exception as e:
                success = False
                logger.error(
                    f"[AWS ERROR] {action.service}:{action.operation} error={str(e)}"
                )

                self.aws.record_operation_failure(
                    episode_id=action.episode_id,
                    task_id=action.task_id,
                    service=action.service,
                    operation=action.operation,
                    error=str(e),
                    error_type=type(e).__name__,
                    principal=action.principal_arn
                    or IAMEnforcer.DEFAULT_ROOT_PRINCIPAL,
                )

            if success:
                self._sync_iam_to_pipeline(action)

            reward, done = self.rewarder.compute_reward(
                pipeline=pipeline, action=action, state=self._state
            )

            self._state.step_count += 1
            result = self.aws.get_last_operation_result()

            logger.info(
                f"[STEP RESULT] step={self._state.step_count}, reward={reward}, done={done}"
            )

            return CloudEnvObservation(
                episode_id=self._state.episode_id,
                task_id=self._state.task_id,
                step_count=self._state.step_count,
                operation_result=result,
                reward=reward,
                done=done,
                current_pipeline_state=pipeline.broken_graph,
            )

    @property
    def state(self) -> State:
        return self._state