from __future__ import annotations

import base64
import json
from pathlib import Path
from time import perf_counter
from typing import Any

from openai import OpenAI

from .prompts import REFINE_SYSTEM_PROMPT, SYSTEM_PROMPT, build_refine_prompt, build_user_prompt
from .schema import ActionDecision


class QwenVisionClient:
    """封装阿里百炼兼容 OpenAI 接口的视觉决策调用。"""

    def __init__(self, model_name: str, base_url: str, api_key: str):
        """初始化千问视觉模型客户端。"""
        self.model_name = model_name
        self.client = OpenAI(api_key=api_key, base_url=base_url)

    def decide_next_action(
        self,
        task: str,
        screenshot_path: str,
        history: list[dict[str, Any]],
        output_schema: dict[str, Any] | None,
    ) -> tuple[ActionDecision, dict[str, Any]]:
        """基于截图和任务上下文生成下一步动作。"""
        prompt = build_user_prompt(task=task, output_schema=output_schema, history=history)
        image_data_url = self._build_data_url(screenshot_path)

        started = perf_counter()
        completion = self.client.chat.completions.create(
            model=self.model_name,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {
                    "role": "user",
                    "content": [
                        {"type": "image_url", "image_url": {"url": image_data_url}},
                        {"type": "text", "text": prompt},
                    ],
                },
            ],
        )
        elapsed_ms = int((perf_counter() - started) * 1000)
        metrics = self._build_metrics(completion, elapsed_ms=elapsed_ms, stage="decide")

        content = completion.choices[0].message.content or "{}"
        payload = self._parse_json(content)

        action = str(payload.get("action", "wait"))
        args = payload.get("args") if isinstance(payload.get("args"), dict) else {}
        reason = str(payload.get("reason", "模型未给出理由"))
        output = payload.get("output") if isinstance(payload.get("output"), dict) else None
        return ActionDecision(action=action, args=args, reason=reason, output=output), metrics

    def refine_output_with_same_image(
        self,
        task: str,
        screenshot_path: str,
        current_output: dict[str, Any],
        output_schema: dict[str, Any] | None,
        max_rounds: int = 4,
    ) -> tuple[dict[str, Any], list[dict[str, Any]]]:
        """基于同一张截图多轮查漏补全，提升列表型结果召回率。"""
        merged_output = current_output
        image_data_url = self._build_data_url(screenshot_path)
        metrics_list: list[dict[str, Any]] = []

        for round_index in range(max_rounds):
            prompt = build_refine_prompt(
                task=task,
                current_output=merged_output,
                output_schema=output_schema,
            )
            started = perf_counter()
            completion = self.client.chat.completions.create(
                model=self.model_name,
                messages=[
                    {"role": "system", "content": REFINE_SYSTEM_PROMPT},
                    {
                        "role": "user",
                        "content": [
                            {"type": "image_url", "image_url": {"url": image_data_url}},
                            {"type": "text", "text": prompt},
                        ],
                    },
                ],
            )
            elapsed_ms = int((perf_counter() - started) * 1000)
            metrics = self._build_metrics(
                completion,
                elapsed_ms=elapsed_ms,
                stage=f"refine-{round_index + 1}",
            )
            metrics_list.append(metrics)

            content = completion.choices[0].message.content or "{}"
            payload = self._parse_json(content)
            candidate = payload.get("output") if isinstance(payload.get("output"), dict) else payload
            if not isinstance(candidate, dict) or not candidate:
                continue

            before = self._count_records(merged_output)
            merged_output = self._merge_outputs(merged_output, candidate)
            after = self._count_records(merged_output)
            if after <= before:
                break

        return merged_output, metrics_list

    def _build_metrics(self, completion: Any, elapsed_ms: int, stage: str) -> dict[str, Any]:
        """提取单次模型请求的耗时和 token 消耗信息。"""
        usage = getattr(completion, "usage", None)
        prompt_tokens = int(getattr(usage, "prompt_tokens", 0) or 0)
        completion_tokens = int(getattr(usage, "completion_tokens", 0) or 0)
        total_tokens = int(getattr(usage, "total_tokens", prompt_tokens + completion_tokens) or 0)
        return {
            "stage": stage,
            "request_id": str(getattr(completion, "id", "")),
            "model": str(getattr(completion, "model", self.model_name)),
            "elapsed_ms": elapsed_ms,
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
            "total_tokens": total_tokens,
        }

    def _merge_outputs(self, base: dict[str, Any], incoming: dict[str, Any]) -> dict[str, Any]:
        """合并两份结果，尽量保留更多信息并对列表去重。"""
        merged = dict(base)
        for key, incoming_value in incoming.items():
            base_value = merged.get(key)
            if isinstance(base_value, dict) and isinstance(incoming_value, dict):
                merged[key] = self._merge_outputs(base_value, incoming_value)
                continue

            if isinstance(base_value, list) and isinstance(incoming_value, list):
                merged[key] = self._merge_list(base_value, incoming_value)
                continue

            if base_value in (None, "", [], {}):
                merged[key] = incoming_value
            elif incoming_value not in (None, "", [], {}):
                merged[key] = incoming_value
        return merged

    def _merge_list(self, first: list[Any], second: list[Any]) -> list[Any]:
        """按内容指纹去重合并列表。"""
        combined = first + second
        seen: set[str] = set()
        result: list[Any] = []
        for item in combined:
            fingerprint = json.dumps(item, ensure_ascii=False, sort_keys=True)
            if fingerprint in seen:
                continue
            seen.add(fingerprint)
            result.append(item)
        return result

    def _count_records(self, payload: dict[str, Any]) -> int:
        """统计结果中的记录规模，用于判断补全是否有增量。"""
        count = 0
        for value in payload.values():
            if isinstance(value, list):
                count += len(value)
            elif isinstance(value, dict):
                count += self._count_records(value)
        return count

    def _build_data_url(self, screenshot_path: str) -> str:
        """将本地截图编码为 data URL 供视觉模型读取。"""
        data = Path(screenshot_path).read_bytes()
        b64 = base64.b64encode(data).decode("utf-8")
        return f"data:image/png;base64,{b64}"

    def _parse_json(self, content: Any) -> dict[str, Any]:
        """尽量稳健地从模型文本中提取 JSON。"""
        if not isinstance(content, str):
            return {}

        text = content.strip()
        if text.startswith("```"):
            text = text.removeprefix("```json").removeprefix("```").removesuffix("```").strip()

        try:
            parsed = json.loads(text)
            return parsed if isinstance(parsed, dict) else {}
        except json.JSONDecodeError:
            start = text.find("{")
            end = text.rfind("}")
            if start >= 0 and end > start:
                try:
                    parsed = json.loads(text[start : end + 1])
                    return parsed if isinstance(parsed, dict) else {}
                except json.JSONDecodeError:
                    return {}
            return {}
