from __future__ import annotations

from datetime import UTC, datetime
from typing import TYPE_CHECKING

import boto3
from aiokafka.errors import KafkaError
from botocore.config import Config as BotoConfig
from fastapi import APIRouter
from loguru import logger
from matyan_api_models.kafka import IngestionMessage
from pydantic import BaseModel

from matyan_frontier.config import SETTINGS
from matyan_frontier.kafka.producer import get_producer

if TYPE_CHECKING:
    from types_boto3_s3.client import S3Client

rest_router = APIRouter(prefix="/rest")

_s3_client = None


def _get_s3_client() -> S3Client:
    global _s3_client  # noqa: PLW0603
    if _s3_client is None:
        _s3_client = boto3.client(
            "s3",
            endpoint_url=SETTINGS.s3_endpoint,
            aws_access_key_id=SETTINGS.s3_access_key,
            aws_secret_access_key=SETTINGS.s3_secret_key,
            config=BotoConfig(signature_version="s3v4"),
            region_name="us-east-1",
        )
        _ensure_bucket(_s3_client, SETTINGS.s3_bucket)
    return _s3_client


def _ensure_bucket(client: S3Client, bucket: str) -> None:
    from botocore.exceptions import ClientError  # noqa: PLC0415

    try:
        client.head_bucket(Bucket=bucket)
    except ClientError:
        client.create_bucket(Bucket=bucket)
        logger.info("Created S3 bucket {!r}", bucket)


class PresignRequest(BaseModel):
    run_id: str
    artifact_path: str
    content_type: str = "application/octet-stream"


class PresignResponse(BaseModel):
    upload_url: str
    s3_key: str


@rest_router.post("/artifacts/presign", response_model=PresignResponse)
async def presign_upload(body: PresignRequest) -> PresignResponse:
    s3_key = f"{body.run_id}/{body.artifact_path}"

    client = _get_s3_client()
    upload_url = client.generate_presigned_url(
        "put_object",
        Params={
            "Bucket": SETTINGS.s3_bucket,
            "Key": s3_key,
            "ContentType": body.content_type,
        },
        ExpiresIn=SETTINGS.s3_presign_expiry,
    )

    message = IngestionMessage(
        type="blob_ref",
        run_id=body.run_id,
        timestamp=datetime.now(UTC),
        payload={
            "s3_key": s3_key,
            "artifact_path": body.artifact_path,
            "content_type": body.content_type,
        },
    )

    try:
        producer = get_producer()
        await producer.publish(message)
    except KafkaError:
        logger.exception("Failed to publish blob_ref", run_id=body.run_id)

    return PresignResponse(upload_url=upload_url, s3_key=s3_key)
