import logging
import os
from typing import Dict, List, Optional, Tuple

from .actions import parse_base_response, parse_internal_response, py_string
from .client import call_openai_compatible
from .history import (
    build_messages,
    dump_debug_messages,
    ensure_empty_think_prefix,
    previous_actions_text,
    update_folding_state,
)
from .images import image_size_from_base64, image_size_from_bytes, process_image
from .prompts import (
    build_base_system_prompt,
    build_base_tools_def,
    build_instruction_prompt,
    build_internal_system_prompt,
    build_internal_tools_def,
)


logger = None


class _QwenBaseAgent:
    """
    Shared implementation for Qwen computer-use agents.

    Characteristics:
    - OpenAI-compatible API only.
    - XML tool-call output format.
    - History truncation by `history_n`.
    - Old screenshot folding by `image_max` / `fold_size`.
    """

    COLLAPSED_SCREENSHOT_TEXT = "This screenshot has been collapsed."

    def __init__(
        self,
        platform: str = "ubuntu",
        model: str = "qwen-vl",
        max_tokens: int = 32768,
        top_p: float = 0.9,
        temperature: float = 0.0,
        action_space: str = "pyautogui",
        observation_type: str = "screenshot",
        history_n: int = 100,
        add_thought_prefix: bool = False,
        coordinate_type: str = "relative",
        api_backend: str = "openai",
        image_max: int = 20,
        fold_size: int = 10,
        collapse_text: Optional[str] = None,
        base_url: Optional[str] = None,
        api_key: Optional[str] = None,
    ):
        self.platform = platform
        self.model = model
        self.max_tokens = max_tokens
        self.top_p = top_p
        self.temperature = temperature
        self.action_space = action_space
        self.observation_type = observation_type
        self.history_n = history_n
        self.add_thought_prefix = add_thought_prefix
        self.coordinate_type = coordinate_type
        self.api_backend = api_backend
        self.image_max = int(image_max)
        self.fold_size = int(fold_size)
        self.collapse_text = collapse_text or self.COLLAPSED_SCREENSHOT_TEXT
        self.base_url = base_url
        self.api_key = api_key

        if action_space != "pyautogui":
            raise ValueError("QwenAgent supports only pyautogui action space")
        if observation_type != "screenshot":
            raise ValueError("QwenAgent supports only screenshot observations")
        if api_backend != "openai":
            raise ValueError("QwenAgent supports only OpenAI-compatible APIs")
        if self.image_max < 1:
            raise ValueError("image_max must be >= 1")
        if self.fold_size < 1:
            raise ValueError("fold_size must be >= 1")

        self.thoughts: List[str] = []
        self.actions: List[str] = []
        self.observations: List[Dict] = []
        self.responses: List[str] = []
        self.screenshots: List[str] = []
        self.folded_prefix_k = 0

    @staticmethod
    def _py_string(text: str) -> str:
        return py_string(text)

    def _build_tools_def(self, processed_width: int, processed_height: int) -> Dict:
        return build_base_tools_def(processed_width, processed_height, self.coordinate_type)

    def _build_system_prompt(self, tools_def: Dict) -> str:
        return build_base_system_prompt(tools_def, self.collapse_text)

    def _response_transform(self, response: str) -> str:
        return response

    def _debug_message_filename(self, step_idx: int) -> str:
        return f"qwen_messages_step_{step_idx}.json"

    def _build_payload(self, messages: List[Dict]) -> Dict:
        return {
            "model": self.model,
            "messages": messages,
            "max_tokens": self.max_tokens,
            "top_p": self.top_p,
            "temperature": self.temperature,
        }

    def _log_prefix(self) -> str:
        return "Qwen"

    def _parse_response(
        self,
        response: str,
        *,
        original_width: int,
        original_height: int,
        processed_width: int,
        processed_height: int,
    ) -> Tuple[str, List[str]]:
        return parse_base_response(
            response,
            coordinate_type=self.coordinate_type,
            original_width=original_width,
            original_height=original_height,
            processed_width=processed_width,
            processed_height=processed_height,
        )

    def predict(self, instruction: str, obs: Dict) -> Tuple[str, List[str]]:
        screenshot_bytes = obs["screenshot"]

        original_width, original_height = image_size_from_bytes(screenshot_bytes)
        processed_b64 = process_image(screenshot_bytes)
        processed_width, processed_height = image_size_from_base64(processed_b64)

        self.screenshots.append(processed_b64)
        total_steps = len(self.screenshots)
        self.folded_prefix_k = update_folding_state(
            total_steps,
            self.folded_prefix_k,
            self.image_max,
            self.fold_size,
        )

        start_step = max(1, total_steps - self.history_n)
        previous_actions_str = previous_actions_text(self.actions, start_step)

        tools_def = self._build_tools_def(processed_width, processed_height)
        system_prompt = self._build_system_prompt(tools_def)
        instruction_prompt = build_instruction_prompt(instruction, previous_actions_str)

        self.observations.append({"screenshot": processed_b64})
        messages = build_messages(
            system_prompt=system_prompt,
            instruction_prompt=instruction_prompt,
            screenshots=self.screenshots,
            responses=self.responses,
            start_step=start_step,
            total_steps=total_steps,
            folded_prefix_k=self.folded_prefix_k,
            collapse_text=self.collapse_text,
            response_transform=self._response_transform,
        )

        step_idx = total_steps - 1
        dump_debug_messages(messages, self._debug_message_filename(step_idx), logger)

        response = self.call_llm(self._build_payload(messages), self.model)

        if logger:
            logger.info("%s Output: %s", self._log_prefix(), response)
        self.responses.append(response or "")

        low_level_instruction, pyautogui_code = self._parse_response(
            response or "",
            original_width=original_width,
            original_height=original_height,
            processed_width=processed_width,
            processed_height=processed_height,
        )

        if logger:
            logger.info("Low level instruction: %s", low_level_instruction)
            logger.info("Pyautogui code: %s", pyautogui_code)

        self.actions.append(low_level_instruction)
        return response or "", pyautogui_code

    def parse_response(
        self,
        response: str,
        original_width: int = None,
        original_height: int = None,
        processed_width: int = None,
        processed_height: int = None,
    ) -> Tuple[str, List[str]]:
        return parse_base_response(
            response,
            coordinate_type=self.coordinate_type,
            original_width=original_width,
            original_height=original_height,
            processed_width=processed_width,
            processed_height=processed_height,
        )

    def call_llm(self, payload: Dict, model: str) -> str:
        return call_openai_compatible(
            payload,
            model,
            base_url=self.base_url,
            api_key=self.api_key,
            default_max_tokens=self.max_tokens,
            default_temperature=self.temperature,
            default_top_p=self.top_p,
            logger=logger,
        )

    def reset(self, _logger=None, *args, **kwargs):
        global logger
        logger = _logger if _logger is not None else logging.getLogger("desktopenv.qwen_agent")
        self.thoughts = []
        self.actions = []
        self.observations = []
        self.responses = []
        self.screenshots = []
        self.folded_prefix_k = 0


class QwenAgent(_QwenBaseAgent):
    COLLAPSED_SCREENSHOT_TEXT = "This screenshot has been collapsed."

    def __init__(
        self,
        *args,
        enable_thinking: bool = False,
        observation_type: str = "screenshot",
        **kwargs,
    ):
        super().__init__(*args, observation_type="screenshot", **kwargs)
        self.enable_thinking = enable_thinking
        if observation_type != "screenshot":
            raise ValueError("QwenAgent supports only screenshot observations")
        self.observation_type = observation_type

    def _build_tools_def(self, processed_width: int, processed_height: int) -> Dict:
        return build_internal_tools_def(processed_width, processed_height, self.coordinate_type)

    def _build_system_prompt(self, tools_def: Dict) -> str:
        return build_internal_system_prompt(tools_def, self.collapse_text)

    def _response_transform(self, response: str) -> str:
        return ensure_empty_think_prefix(response)

    def _debug_message_filename(self, step_idx: int) -> str:
        return f"qwen_messages_step_{step_idx}.json"

    def _build_payload(self, messages: List[Dict]) -> Dict:
        payload = super()._build_payload(messages)
        base_url = self.base_url or os.environ.get("OPENAI_BASE_URL", "")
        if "dashscope" in base_url.lower():
            extra_body = dict(payload.get("extra_body") or {})
            extra_body["enable_thinking"] = bool(self.enable_thinking)
            payload["extra_body"] = extra_body
        return payload

    def _log_prefix(self) -> str:
        return "Qwen"

    def _parse_response(
        self,
        response: str,
        *,
        original_width: int,
        original_height: int,
        processed_width: int,
        processed_height: int,
    ) -> Tuple[str, List[str]]:
        return parse_internal_response(
            response,
            coordinate_type=self.coordinate_type,
            original_width=original_width,
            original_height=original_height,
            processed_width=processed_width,
            processed_height=processed_height,
        )

    def parse_response(
        self,
        response: str,
        original_width: int = None,
        original_height: int = None,
        processed_width: int = None,
        processed_height: int = None,
    ) -> Tuple[str, List[str]]:
        return parse_internal_response(
            response,
            coordinate_type=self.coordinate_type,
            original_width=original_width,
            original_height=original_height,
            processed_width=processed_width,
            processed_height=processed_height,
        )
