"""Short-lived uploaded images consumed by the realtime user-input protocol."""

from __future__ import annotations

import base64
import hashlib
import hmac
import json
import time
from dataclasses import dataclass
from io import BytesIO
from uuid import uuid4

from PIL import Image, ImageOps, UnidentifiedImageError

import settings
from services.object_storage import delete_object, get_object, put_object


INPUT_ASSET_PREFIX = "conversation-inputs/"


@dataclass(frozen=True)
class StoredInputImage:
    token: str
    width: int
    height: int
    size_bytes: int


def _b64encode(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")


def _b64decode(data: str) -> bytes:
    padding = "=" * (-len(data) % 4)
    return base64.urlsafe_b64decode(f"{data}{padding}".encode("ascii"))


def _signature(payload: str) -> str:
    digest = hmac.new(
        settings.AUTH_SECRET_KEY.encode("utf-8"),
        payload.encode("ascii"),
        hashlib.sha256,
    ).digest()
    return _b64encode(digest)


def _normalize_image(data: bytes) -> tuple[bytes, int, int]:
    try:
        source = Image.open(BytesIO(data))
        width, height = source.size
        if width <= 0 or height <= 0:
            raise ValueError("图片尺寸无效")
        if width * height > settings.INPUT_IMAGE_MAX_PIXELS:
            raise ValueError("图片像素尺寸过大")
        source.load()
    except (UnidentifiedImageError, OSError) as exc:
        raise ValueError("文件不是可识别的图片") from exc

    image = ImageOps.exif_transpose(source)
    image.thumbnail(
        (settings.INPUT_IMAGE_MAX_EDGE, settings.INPUT_IMAGE_MAX_EDGE),
        Image.Resampling.LANCZOS,
    )
    if image.mode in {"RGBA", "LA"} or "transparency" in image.info:
        rgba = image.convert("RGBA")
        background = Image.new("RGB", rgba.size, "white")
        background.paste(rgba, mask=rgba.getchannel("A"))
        image = background
    else:
        image = image.convert("RGB")

    buffer = BytesIO()
    image.save(buffer, format="JPEG", quality=85, optimize=True)
    return buffer.getvalue(), image.width, image.height


def store_input_image(data: bytes) -> StoredInputImage:
    normalized, width, height = _normalize_image(data)
    key = f"{INPUT_ASSET_PREFIX}{uuid4().hex}.jpg"
    put_object(key, normalized, "image/jpeg")
    expires_at = int(time.time()) + settings.INPUT_IMAGE_TOKEN_TTL_SECONDS
    encoded_payload = _b64encode(
        json.dumps(
            {"key": key, "exp": expires_at},
            separators=(",", ":"),
        ).encode("utf-8")
    )
    token = f"{encoded_payload}.{_signature(encoded_payload)}"
    return StoredInputImage(
        token=token,
        width=width,
        height=height,
        size_bytes=len(normalized),
    )


def _key_from_token(token: str) -> str:
    if not token or "." not in token:
        raise ValueError("图片附件 token 无效")
    encoded_payload, signature = token.rsplit(".", 1)
    if not hmac.compare_digest(_signature(encoded_payload), signature):
        raise ValueError("图片附件 token 签名无效")
    try:
        payload = json.loads(_b64decode(encoded_payload))
        key = str(payload.get("key") or "")
        expires_at = int(payload.get("exp") or 0)
    except (ValueError, TypeError, json.JSONDecodeError) as exc:
        raise ValueError("图片附件 token 内容无效") from exc
    if expires_at < int(time.time()):
        raise ValueError("图片附件已过期，请重新添加")
    if not key.startswith(INPUT_ASSET_PREFIX):
        raise ValueError("图片附件存储位置无效")
    return key


def consume_input_image(token: str) -> bytes:
    """Read a signed image once and remove its temporary object."""

    key = _key_from_token(token)
    data = get_object(key)
    try:
        delete_object(key)
    except Exception:
        # Consumption succeeded. A stale temporary object must not break the turn.
        pass
    return data


def discard_input_image(token: str) -> None:
    delete_object(_key_from_token(token))