"""企业微信 API 模式：上传图片并回复。"""

from __future__ import annotations

import base64
import hashlib
import logging
from pathlib import Path
from typing import Any

from aibot import generate_req_id

logger = logging.getLogger(__name__)

CHUNK_SIZE = 512 * 1024
MAX_IMAGE_BYTES = 9 * 1024 * 1024


def _ensure_image_size(path: Path) -> bytes:
    data = path.read_bytes()
    if len(data) > MAX_IMAGE_BYTES:
        raise RuntimeError(
            f"截图过大（{len(data) // 1024}KB），请缩小页面或使用 viewport 截图（上限 9MB）"
        )
    return data


def _response_body(frame: dict[str, Any]) -> dict[str, Any]:
    if frame.get("errcode", 0) != 0:
        raise RuntimeError(
            f"企微接口错误 errcode={frame.get('errcode')} errmsg={frame.get('errmsg')}"
        )
    body = frame.get("body")
    return body if isinstance(body, dict) else {}


async def upload_image(ws_client: Any, image_path: str | Path) -> str:
    path = Path(image_path)
    if not path.exists():
        raise RuntimeError(f"截图不存在: {path}")

    data = _ensure_image_size(path)
    md5 = hashlib.md5(data).hexdigest()
    chunks = [data[i : i + CHUNK_SIZE] for i in range(0, len(data), CHUNK_SIZE)]
    total_chunks = len(chunks)

    manager = ws_client._ws_manager

    init_frame = await manager.send_reply(
        generate_req_id("upload_init"),
        {
            "type": "image",
            "filename": path.name,
            "total_size": len(data),
            "total_chunks": total_chunks,
            "md5": md5,
        },
        "aibot_upload_media_init",
    )
    upload_id = _response_body(init_frame).get("upload_id")
    if not upload_id:
        raise RuntimeError("上传初始化失败：未返回 upload_id")

    for index, chunk in enumerate(chunks):
        chunk_frame = await manager.send_reply(
            generate_req_id("upload_chunk"),
            {
                "upload_id": upload_id,
                "chunk_index": index,
                "base64_data": base64.b64encode(chunk).decode("ascii"),
            },
            "aibot_upload_media_chunk",
        )
        _response_body(chunk_frame)

    finish_frame = await manager.send_reply(
        generate_req_id("upload_finish"),
        {"upload_id": upload_id},
        "aibot_upload_media_finish",
    )
    media_id = _response_body(finish_frame).get("media_id")
    if not media_id:
        raise RuntimeError("上传完成但未返回 media_id")

    logger.info("图片已上传 media_id=%s…", str(media_id)[:12])
    return str(media_id)


async def reply_image(ws_client: Any, frame: dict[str, Any], media_id: str) -> None:
    await ws_client.reply(
        frame,
        {
            "msgtype": "image",
            "image": {"media_id": media_id},
        },
    )