"""image2ppt 企业 API 的官方 Python 客户端。

把一批图片 / PDF 提交给 image2ppt，等它转换成一个可编辑的 PPTX，再下载下来。
整个过程是异步的：提交后拿到一个任务号，轮询任务状态，完成后下载成品。

上传前，客户端会先把图片压到服务器同样的规格（≤2000px、≤1MB、JPEG），这样
服务器那趟压缩就是空操作（passthrough），省掉一次重复计算、也少传字节。PDF
不在客户端处理，原样上传交由服务器渲染（渲染质量由服务器统一把控）。

依赖：`requests` + `Pillow`（`uv add requests pillow`），Python 3.9 及以上。

最小用法::

    from image2ppt_client import Image2PPTClient

    client = Image2PPTClient(api_key="i2p_live_你的密钥")

    # 一步到位：提交 → 等待 → 下载
    job = client.convert(["1.png", "2.png", "report.pdf"], dest_path="out.pptx")
    print("用掉积分：", job.credits_used, "退回积分：", job.credits_refunded)

    # 或者分步来，自己控制轮询
    job = client.submit(["1.png"], locale="zh-CN", aspect_ratio="16:9")
    job = client.wait(job.job_id)
    client.download(job.job_id, "out.pptx")

    print("剩余积分：", client.account()["credits"])
"""

from __future__ import annotations

import io
import mimetypes
import os
import time
from dataclasses import dataclass
from typing import Any, Dict, Optional, Sequence

import requests
from PIL import Image

# ------ 客户端图片压缩（与服务器保持一致）--------------------------------- #
# 服务器每个上传入口最终都会跑一遍同样的压缩（src/services/image-loader.ts 的
# compressImageForUpload）。客户端先压好，服务器那趟就变成 passthrough。以下
# 数值必须与服务器、以及 web/assets/js/upload.ts、canvas-compress.ts 严格同步
# ——改一处要四处一起改，否则要么没压够、要么服务器再压一遍白费功夫。
_UPLOAD_TARGET_BYTES = 1024 * 1024
_UPLOAD_MAX_DIM = 2000
_UPLOAD_QUALITY_LADDER = (90, 85, 80)
# 只有 PNG / JPEG 能原样放行；WebP / GIF 即便够小也转码成 JPEG（对齐服务器：它
# 对非 PNG/JPEG 一律先转再判断）。
_PASSTHROUGH_MIMES = frozenset({"image/png", "image/jpeg"})
_IMAGE_MIMES = frozenset({"image/png", "image/jpeg", "image/webp", "image/gif"})


def _compress_image_for_upload(raw: bytes, mime: str) -> tuple[bytes, str]:
    """把一张图片压到与服务器一致的规格，返回 (bytes, mime)。

    规则逐条对齐服务器 compressImageForUpload：
      - PNG/JPEG 且长边 ≤2000px、字节 ≤1MB → 原样返回（passthrough）。
      - 否则：等比缩进 2000×2000（只缩不放）、白底去透明、JPEG 质量 90→85→80
        阶梯，压到 ≤1MB 或质量触底为止。
      - 兜底：压完反而比原图大（本就低质量的源图会这样）→ 退回原图，绝不「更糊
        还更大」。

    只处理图片；PDF 不走这里，原样上传交服务器渲染。
    """
    with Image.open(io.BytesIO(raw)) as img:
        img.load()  # GIF / WebP 动图只取第一帧（Pillow 默认帧）
        width, height = img.size
        within_budget = (
            len(raw) <= _UPLOAD_TARGET_BYTES and max(width, height) <= _UPLOAD_MAX_DIM
        )
        if within_budget and mime in _PASSTHROUGH_MIMES:
            return raw, mime

        scaled = img.copy()
        # thumbnail 等比缩且不放大，等价服务器的 fit:inside + withoutEnlargement。
        if max(width, height) > _UPLOAD_MAX_DIM:
            scaled.thumbnail((_UPLOAD_MAX_DIM, _UPLOAD_MAX_DIM), Image.LANCZOS)

        # 白底 flatten，去掉透明通道（对齐服务器 .flatten({background:'#ffffff'})）。
        has_alpha = scaled.mode in ("RGBA", "LA") or (
            scaled.mode == "P" and "transparency" in scaled.info
        )
        if has_alpha:
            rgba = scaled.convert("RGBA")
            flattened = Image.new("RGB", rgba.size, (255, 255, 255))
            flattened.paste(rgba, mask=rgba.split()[-1])
            scaled = flattened
        else:
            scaled = scaled.convert("RGB")

        compressed = None
        for quality in _UPLOAD_QUALITY_LADDER:
            buffer = io.BytesIO()
            scaled.save(buffer, format="JPEG", quality=quality)
            compressed = buffer.getvalue()
            if len(compressed) <= _UPLOAD_TARGET_BYTES:
                break

    if compressed is None or len(compressed) >= len(raw):
        return raw, mime
    return compressed, "image/jpeg"


__all__ = [
    "Image2PPTClient",
    "Job",
    "Image2PPTError",
    "AuthenticationError",
    "InvalidFileError",
    "TooManySlidesError",
    "InsufficientCreditsError",
    "RateLimitedError",
    "JobNotFoundError",
    "NotReadyError",
    "JobFailedError",
    "Image2PPTTimeoutError",
]

DEFAULT_BASE_URL = "https://image2ppt.com"


# --------------------------------------------------------------------------- #
# 异常
# --------------------------------------------------------------------------- #
class Image2PPTError(Exception):
    """所有客户端错误的基类。

    携带服务端返回的 HTTP 状态码、错误码（错误信封里的 ``error.code``）和文字说明，
    方便调用方按 ``code`` 做分支处理。
    """

    def __init__(
        self,
        message: str,
        *,
        status_code: Optional[int] = None,
        code: Optional[str] = None,
    ) -> None:
        super().__init__(message)
        self.message = message
        self.status_code = status_code
        self.code = code

    def __str__(self) -> str:
        parts = []
        if self.status_code is not None:
            parts.append(f"HTTP {self.status_code}")
        if self.code:
            parts.append(self.code)
        prefix = " ".join(parts)
        return f"[{prefix}] {self.message}" if prefix else self.message


class AuthenticationError(Image2PPTError):
    """API 密钥无效或缺失（401 INVALID_API_KEY）。"""


class InvalidFileError(Image2PPTError):
    """上传的文件不合法：格式不支持或单文件超过 35MB（400 INVALID_FILE）。"""


class TooManySlidesError(Image2PPTError):
    """一次提交的总页数超过 50 页上限（400 TOO_MANY_SLIDES）。"""


class InsufficientCreditsError(Image2PPTError):
    """账户可用积分不足以覆盖这次提交（402 INSUFFICIENT_CREDITS）。"""


class RateLimitedError(Image2PPTError):
    """触发限流（429 RATE_LIMITED）。

    ``retry_after`` 是服务端建议的等待秒数（来自 Retry-After 响应头），
    过了这段时间再重试即可。
    """

    def __init__(
        self,
        message: str,
        *,
        status_code: Optional[int] = None,
        code: Optional[str] = None,
        retry_after: Optional[float] = None,
    ) -> None:
        super().__init__(message, status_code=status_code, code=code)
        self.retry_after = retry_after


class JobNotFoundError(Image2PPTError):
    """任务号不存在，或不属于当前密钥所在账户（404 JOB_NOT_FOUND）。"""


class NotReadyError(Image2PPTError):
    """任务还没转换完成，成品尚不可下载（409 NOT_READY）。"""


class JobFailedError(Image2PPTError):
    """任务最终以失败告终（由 ``wait`` 在轮询到 status=failed 时抛出）。

    ``job`` 是失败时的任务快照，``code`` / ``message`` 来自任务的 ``error`` 字段。
    """

    def __init__(
        self,
        message: str,
        *,
        code: Optional[str] = None,
        job: Optional["Job"] = None,
    ) -> None:
        super().__init__(message, code=code)
        self.job = job


class Image2PPTTimeoutError(Image2PPTError):
    """轮询等待超过了给定的 ``timeout`` 上限，任务仍未结束。

    这不代表任务失败——它可能还在跑。可以稍后拿 ``job_id`` 重新 ``wait`` 或查询状态。
    （名字带前缀是为了不遮蔽 Python 内置的 ``TimeoutError``。）
    """

    def __init__(self, message: str, *, job_id: Optional[str] = None) -> None:
        super().__init__(message)
        self.job_id = job_id


# 服务端错误码 → 异常类。未列出的错误码回退到状态码映射，再回退到基类。
_CODE_TO_EXC: Dict[str, type] = {
    "INVALID_API_KEY": AuthenticationError,
    "INVALID_FILE": InvalidFileError,
    "TOO_MANY_SLIDES": TooManySlidesError,
    "INSUFFICIENT_CREDITS": InsufficientCreditsError,
    "RATE_LIMITED": RateLimitedError,
    "JOB_NOT_FOUND": JobNotFoundError,
    "NOT_READY": NotReadyError,
}
_STATUS_TO_EXC: Dict[int, type] = {
    401: AuthenticationError,
    402: InsufficientCreditsError,
    404: JobNotFoundError,
    409: NotReadyError,
    429: RateLimitedError,
}


# --------------------------------------------------------------------------- #
# 任务数据模型
# --------------------------------------------------------------------------- #
@dataclass
class Job:
    """一个转换任务的状态快照。

    字段随来源不同而有的为 None：``submit`` 的响应只带 ``credits_reserved``，
    ``get_job`` 的响应带 ``credits_used`` / ``credits_refunded`` / ``download_url`` 等。
    """

    job_id: str
    status: str  # pending | processing | completed | failed
    slide_count: Optional[int] = None
    progress: Optional[int] = None  # 0-100
    credits_reserved: Optional[int] = None  # 仅提交响应：锁定的积分
    credits_used: Optional[int] = None  # 结算后实际扣除
    credits_refunded: Optional[int] = None  # 部分成功时退回的失败页积分
    created_at: Optional[str] = None
    completed_at: Optional[str] = None
    download_url: Optional[str] = None  # 仅 completed，相对路径
    error: Optional[Dict[str, Any]] = None  # 仅 failed，{code, message}
    raw: Optional[Dict[str, Any]] = None  # 原始响应体，便于访问新增字段

    @property
    def is_completed(self) -> bool:
        """任务是否已成功完成（成品可下载）。"""
        return self.status == "completed"

    @property
    def is_failed(self) -> bool:
        """任务是否已失败。"""
        return self.status == "failed"

    @property
    def is_terminal(self) -> bool:
        """任务是否已到终态（成功或失败），不会再变化。"""
        return self.status in ("completed", "failed")

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "Job":
        """从服务端 JSON 构造 Job；兼容提交响应和状态查询响应两种形状。"""
        return cls(
            job_id=data["jobId"],
            status=data["status"],
            slide_count=data.get("slideCount"),
            progress=data.get("progress"),
            credits_reserved=data.get("creditsReserved"),
            credits_used=data.get("creditsUsed"),
            credits_refunded=data.get("creditsRefunded"),
            created_at=data.get("createdAt"),
            completed_at=data.get("completedAt"),
            download_url=data.get("downloadUrl"),
            error=data.get("error"),
            raw=data,
        )


# --------------------------------------------------------------------------- #
# 客户端
# --------------------------------------------------------------------------- #
class Image2PPTClient:
    """image2ppt 企业 API 客户端。

    参数：
        api_key: 企业 API 密钥（形如 ``i2p_live_...``），在账户页自助创建。
        base_url: 服务地址，默认 ``https://image2ppt.com``。
        timeout: 单次 HTTP 请求的超时秒数（不是整个任务的等待上限）。
        session: 可选，注入自定义的 ``requests.Session``（便于测试或复用连接池）。
    """

    #: 支持的输入文件扩展名 → MIME 类型（用于 multipart 上传时标注）。
    _MIME_BY_EXT = {
        ".png": "image/png",
        ".jpg": "image/jpeg",
        ".jpeg": "image/jpeg",
        ".webp": "image/webp",
        ".gif": "image/gif",
        ".pdf": "application/pdf",
    }

    def __init__(
        self,
        api_key: str,
        base_url: str = DEFAULT_BASE_URL,
        *,
        timeout: float = 60.0,
        session: Optional[requests.Session] = None,
    ) -> None:
        if not api_key:
            raise ValueError("api_key 不能为空")
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout
        self._session = session or requests.Session()
        self._session.headers.update({"Authorization": f"Bearer {api_key}"})

    # ----- 公开方法 ----------------------------------------------------- #
    def submit(
        self,
        paths: Sequence[str],
        *,
        locale: Optional[str] = None,
        aspect_ratio: Optional[str] = None,
    ) -> Job:
        """提交一批文件，创建一个转换任务。

        参数：
            paths: 本地文件路径列表（1 个或多个）。支持 png/jpeg/webp/gif/pdf，
                单文件不超过 35MB。图片按 1 页算，PDF 按实际页数算，总页数不超过 50。
            locale: 可选，``zh-CN`` 或 ``en``，默认 ``zh-CN``。
            aspect_ratio: 可选，``auto`` / ``16:9`` / ``4:3``，默认 ``auto``。

        返回：一个 ``Job``，其 ``status`` 为 ``pending``，并带上 ``slide_count`` 和
        ``credits_reserved``（提交时锁定的积分）。

        失败会抛出对应异常：AuthenticationError / InvalidFileError /
        TooManySlidesError / InsufficientCreditsError / RateLimitedError。
        """
        paths = list(paths)
        if not paths:
            raise ValueError("至少要提交一个文件")

        data: Dict[str, str] = {}
        if locale is not None:
            data["locale"] = locale
        if aspect_ratio is not None:
            data["aspectRatio"] = aspect_ratio

        opened = []
        multipart = []
        try:
            for path in paths:
                filename = os.path.basename(path)
                mime = self._guess_mime(filename)
                if mime in _IMAGE_MIMES:
                    # 图片：客户端先压到服务器规格，服务器那趟就是 passthrough。
                    with open(path, "rb") as fh:
                        raw = fh.read()
                    payload, out_mime = _compress_image_for_upload(raw, mime)
                    if out_mime == "image/jpeg" and not filename.lower().endswith(
                        (".jpg", ".jpeg")
                    ):
                        # 压成了 JPEG，文件名扩展名随之对齐，免得扩展名与内容不符。
                        filename = os.path.splitext(filename)[0] + ".jpg"
                    multipart.append(("files", (filename, payload, out_mime)))
                else:
                    # PDF 等非图片：原样上传（流式），渲染留在服务器。
                    handle = open(path, "rb")
                    opened.append(handle)
                    multipart.append(("files", (filename, handle, mime)))
            resp = self._session.post(
                f"{self.base_url}/api/v1/jobs",
                files=multipart,
                data=data,
                timeout=self.timeout,
            )
        finally:
            for handle in opened:
                handle.close()

        return Job.from_dict(self._parse_json(resp))

    def get_job(self, job_id: str) -> Job:
        """查询任务当前状态，返回一个 ``Job`` 快照。任务不存在会抛 JobNotFoundError。"""
        resp = self._session.get(
            f"{self.base_url}/api/v1/jobs/{job_id}",
            timeout=self.timeout,
        )
        return Job.from_dict(self._parse_json(resp))

    def wait(
        self,
        job_id: str,
        *,
        poll_interval: float = 5.0,
        timeout: float = 1800.0,
    ) -> Job:
        """轮询直到任务结束，返回成功完成的 ``Job``。

        轮询间隔从 ``poll_interval`` 起，逐步退避到最多 15 秒。遇到限流（429）会
        按 Retry-After 建议的秒数等待后再继续。任务失败会抛 JobFailedError；
        超过 ``timeout`` 秒仍未结束会抛 Image2PPTTimeoutError（任务本身可能还在跑）。

        参数：
            job_id: 任务号。
            poll_interval: 起始轮询间隔秒数（默认 5）。
            timeout: 整体等待上限秒数（默认 1800，即 30 分钟）。
        """
        deadline = time.monotonic() + timeout
        interval = poll_interval
        while True:
            try:
                job = self.get_job(job_id)
            except RateLimitedError as exc:
                sleep_for = exc.retry_after if exc.retry_after is not None else interval
                self._sleep_until(deadline, sleep_for, job_id)
                continue

            if job.is_completed:
                return job
            if job.is_failed:
                err = job.error or {}
                raise JobFailedError(
                    err.get("message") or "转换失败",
                    code=err.get("code"),
                    job=job,
                )

            self._sleep_until(deadline, interval, job_id)
            interval = min(interval * 1.5, 15.0)

    def download(self, job_id: str, dest_path: str) -> str:
        """把已完成任务的成品 PPTX 流式写到 ``dest_path``，返回该路径。

        任务未完成会抛 NotReadyError（409）；任务不存在会抛 JobNotFoundError（404）。
        """
        resp = self._session.get(
            f"{self.base_url}/api/v1/jobs/{job_id}/download",
            stream=True,
            timeout=self.timeout,
        )
        try:
            if not resp.ok:
                self._raise_for_error(resp)
            with open(dest_path, "wb") as out:
                for chunk in resp.iter_content(chunk_size=65536):
                    if chunk:
                        out.write(chunk)
        finally:
            resp.close()
        return dest_path

    def convert(
        self,
        paths: Sequence[str],
        dest_path: str,
        *,
        locale: Optional[str] = None,
        aspect_ratio: Optional[str] = None,
        poll_interval: float = 5.0,
        timeout: float = 1800.0,
    ) -> Job:
        """一条龙：提交 → 等待完成 → 下载到 ``dest_path``，返回完成时的 ``Job``。

        参数含义同 ``submit`` 与 ``wait``。适合"给我一批图、还我一个 PPTX"的同步场景。
        """
        job = self.submit(paths, locale=locale, aspect_ratio=aspect_ratio)
        completed = self.wait(job.job_id, poll_interval=poll_interval, timeout=timeout)
        self.download(completed.job_id, dest_path)
        return completed

    def account(self) -> Dict[str, Any]:
        """查询账户信息，返回 ``{"email": ..., "credits": 可用积分}``。"""
        resp = self._session.get(
            f"{self.base_url}/api/v1/account",
            timeout=self.timeout,
        )
        return self._parse_json(resp)

    # ----- 内部辅助 ----------------------------------------------------- #
    def _guess_mime(self, filename: str) -> str:
        ext = os.path.splitext(filename)[1].lower()
        if ext in self._MIME_BY_EXT:
            return self._MIME_BY_EXT[ext]
        guessed, _ = mimetypes.guess_type(filename)
        return guessed or "application/octet-stream"

    def _sleep_until(self, deadline: float, seconds: float, job_id: str) -> None:
        """睡 ``seconds`` 秒，但不越过 ``deadline``；越过则抛 TimeoutError。"""
        remaining = deadline - time.monotonic()
        if remaining <= 0:
            raise Image2PPTTimeoutError(f"等待任务 {job_id} 超时", job_id=job_id)
        time.sleep(min(seconds, remaining))

    def _parse_json(self, resp: requests.Response) -> Dict[str, Any]:
        """2xx 返回 JSON 体；否则按错误信封抛出对应异常。"""
        if not resp.ok:
            self._raise_for_error(resp)
        return resp.json()

    def _raise_for_error(self, resp: requests.Response) -> None:
        """解析统一错误信封 ``{"error": {code, message}}`` 并抛出映射后的异常。"""
        code: Optional[str] = None
        message: Optional[str] = None
        try:
            body = resp.json()
            err = body.get("error") if isinstance(body, dict) else None
            if isinstance(err, dict):
                code = err.get("code")
                message = err.get("message")
        except ValueError:
            pass  # 错误响应不是 JSON（如网关返回的 HTML），退回状态码文案
        message = message or f"请求失败（HTTP {resp.status_code}）"

        if resp.status_code == 429:
            raise RateLimitedError(
                message,
                status_code=429,
                code=code or "RATE_LIMITED",
                retry_after=self._parse_retry_after(resp.headers.get("Retry-After")),
            )

        exc_cls = _CODE_TO_EXC.get(code or "") or _STATUS_TO_EXC.get(
            resp.status_code, Image2PPTError
        )
        raise exc_cls(message, status_code=resp.status_code, code=code)

    @staticmethod
    def _parse_retry_after(value: Optional[str]) -> Optional[float]:
        """Retry-After 头按秒数解析（契约规定为整数秒）；无法解析则返回 None。"""
        if not value:
            return None
        try:
            return float(value)
        except (TypeError, ValueError):
            return None