from __future__ import annotations

import io
import uuid
from functools import lru_cache

import boto3
from botocore.client import BaseClient
from botocore.exceptions import ClientError
from PIL import Image, UnidentifiedImageError

from .config import settings


class ScreenshotError(ValueError):
    pass


ALLOWED_FORMATS = {"JPEG": ("jpg", "image/jpeg"), "PNG": ("png", "image/png"), "WEBP": ("webp", "image/webp")}


@lru_cache
def client() -> BaseClient:
    return boto3.client("s3", endpoint_url=settings.s3_endpoint_url, aws_access_key_id=settings.s3_access_key, aws_secret_access_key=settings.s3_secret_key, region_name="us-east-1")


@lru_cache
def public_client() -> BaseClient:
    return boto3.client("s3", endpoint_url=settings.s3_public_endpoint_url, aws_access_key_id=settings.s3_access_key, aws_secret_access_key=settings.s3_secret_key, region_name="us-east-1")


def prepare_image(raw: bytes) -> tuple[bytes, str, str]:
    if not raw or len(raw) > settings.screenshot_max_bytes:
        raise ScreenshotError("screenshot must be between 1 byte and 8 MB")
    try:
        with Image.open(io.BytesIO(raw)) as source:
            source.verify()
        with Image.open(io.BytesIO(raw)) as source:
            fmt = (source.format or "").upper()
            if fmt not in ALLOWED_FORMATS:
                raise ScreenshotError("only JPEG, PNG and WebP images are accepted")
            if source.width * source.height > 40_000_000:
                raise ScreenshotError("image dimensions are too large")
            extension, mime = ALLOWED_FORMATS[fmt]
            clean = source.convert("RGB") if fmt in {"JPEG", "WEBP"} else source.convert("RGBA")
            output = io.BytesIO()
            clean.save(output, format=fmt, optimize=True)
            return output.getvalue(), extension, mime
    except (UnidentifiedImageError, OSError) as exc:
        raise ScreenshotError("file is not a valid image") from exc


def upload_screenshot(raw: bytes) -> str:
    body, extension, mime = prepare_image(raw)
    key = f"reports/{uuid.uuid4()}.{extension}"
    s3 = client()
    try:
        s3.head_bucket(Bucket=settings.s3_bucket)
    except ClientError:
        s3.create_bucket(Bucket=settings.s3_bucket)
    s3.put_object(Bucket=settings.s3_bucket, Key=key, Body=body, ContentType=mime)
    return key


def signed_screenshot_url(key: str) -> str:
    return public_client().generate_presigned_url("get_object", Params={"Bucket": settings.s3_bucket, "Key": key}, ExpiresIn=900)


def delete_screenshot(key: str) -> None:
    client().delete_object(Bucket=settings.s3_bucket, Key=key)