from __future__ import annotations

import io
import uuid
from functools import lru_cache

import boto3
from botocore.client import BaseClient
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")}
ALLOWED_UPLOADS = {"image/jpeg": {".jpg", ".jpeg"}, "image/png": {".png"}, "image/webp": {".webp"}}


def validate_upload_metadata(filename: str | None, content_type: str | None) -> None:
    mime = (content_type or "").lower()
    suffix = f".{(filename or '').rsplit('.', 1)[-1].lower()}" if "." in (filename or "") else ""
    if mime not in ALLOWED_UPLOADS or suffix not in ALLOWED_UPLOADS[mime]:
        raise ScreenshotError("filename extension and MIME type must match JPEG, PNG or 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, *, filename: str | None = None, content_type: str | None = None) -> str:
    validate_upload_metadata(filename, content_type)
    body, extension, mime = prepare_image(raw)
    key = f"reports/{uuid.uuid4()}.{extension}"
    s3 = client()
    s3.head_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)