"""GCS-backed encrypted blob storage for scraper inputs and results."""

from __future__ import annotations

import asyncio
import base64
import json
import logging
from datetime import datetime, timedelta, timezone
from typing import Any, Dict, List, Optional, TYPE_CHECKING
from uuid import UUID

from .kms_envelope import EncryptedSecret, get_envelope_encryption
from .. import settings

logger = logging.getLogger(__name__)

if TYPE_CHECKING:
    pass

_storage_module: Optional[Any] = None
_storage_client: Optional[Any] = None


def _get_storage_module():
    global _storage_module
    if _storage_module is None:
        from google.cloud import storage
        _storage_module = storage
    return _storage_module


def _get_storage_client():
    global _storage_client
    if _storage_client is None:
        storage = _get_storage_module()
        _storage_client = storage.Client()
    return _storage_client


def _encode_encrypted(encrypted: EncryptedSecret) -> bytes:
    payload = {
        "encrypted_value": base64.b64encode(encrypted.encrypted_value).decode("ascii"),
        "encryption_nonce": base64.b64encode(encrypted.encryption_nonce).decode("ascii"),
        "encrypted_dek": base64.b64encode(encrypted.encrypted_dek).decode("ascii"),
        "kms_key_version": encrypted.kms_key_version,
    }
    return json.dumps(payload).encode("utf-8")


def _decode_encrypted(payload: bytes) -> EncryptedSecret:
    data = json.loads(payload.decode("utf-8"))
    return EncryptedSecret(
        encrypted_value=base64.b64decode(data["encrypted_value"]),
        encryption_nonce=base64.b64decode(data["encryption_nonce"]),
        encrypted_dek=base64.b64decode(data["encrypted_dek"]),
        kms_key_version=data["kms_key_version"],
    )


def _build_blob_path(workspace_id: str, job_id: UUID, name: str) -> str:
    return f"{workspace_id}/{job_id}/{name}.json.enc"


def _build_gs_uri(bucket: str, path: str) -> str:
    return f"gs://{bucket}/{path}"


class ScraperBlobStorage:
    def __init__(self, bucket_name: Optional[str] = None):
        self.bucket_name = bucket_name or settings.SCRAPER_DATA_BUCKET
        self.encryption = get_envelope_encryption()

    async def store_inputs(self, workspace_id: str, job_id: UUID, rows: List[Dict[str, Any]]) -> str:
        path = _build_blob_path(workspace_id, job_id, "inputs")
        await asyncio.to_thread(self._upload_payload, path, rows)
        return _build_gs_uri(self.bucket_name, path)

    async def store_results(self, workspace_id: str, job_id: UUID, results: List[Dict[str, Any]]) -> str:
        path = _build_blob_path(workspace_id, job_id, "results")
        await asyncio.to_thread(self._upload_payload, path, results)
        return _build_gs_uri(self.bucket_name, path)

    async def get_inputs(self, workspace_id: str, job_id: UUID) -> List[Dict[str, Any]]:
        path = _build_blob_path(workspace_id, job_id, "inputs")
        return await asyncio.to_thread(self._download_payload, path)

    async def get_results(self, workspace_id: str, job_id: UUID) -> List[Dict[str, Any]]:
        path = _build_blob_path(workspace_id, job_id, "results")
        return await asyncio.to_thread(self._download_payload, path)

    async def get_results_chunk(
        self,
        workspace_id: str,
        job_id: UUID,
        cursor: int,
        limit: int,
    ) -> Dict[str, Any]:
        results = await self.get_results(workspace_id, job_id)
        if cursor < 0:
            cursor = 0
        if limit <= 0:
            limit = 100
        next_cursor = cursor + limit
        chunk = results[cursor:next_cursor]
        has_more = next_cursor < len(results)
        return {
            "results": chunk,
            "next_cursor": next_cursor,
            "has_more": has_more,
        }

    async def delete_blobs(self, workspace_id: str, job_id: UUID) -> None:
        await asyncio.to_thread(self._delete_blob, _build_blob_path(workspace_id, job_id, "inputs"))
        await asyncio.to_thread(self._delete_blob, _build_blob_path(workspace_id, job_id, "results"))

    def _upload_payload(self, path: str, payload: Any) -> None:
        bucket = _get_storage_client().bucket(self.bucket_name)
        blob = bucket.blob(path)
        plaintext = json.dumps(payload, default=str)
        encrypted = self.encryption.encrypt(plaintext)
        blob.upload_from_string(_encode_encrypted(encrypted), content_type="application/json")
        expires_at = datetime.now(timezone.utc) + timedelta(seconds=settings.SCRAPER_DATA_TTL_SECONDS)
        blob.metadata = {
            "expires_at": expires_at.isoformat(),
            "kms_key_version": encrypted.kms_key_version,
        }
        blob.patch()

    def _download_payload(self, path: str) -> Any:
        bucket = _get_storage_client().bucket(self.bucket_name)
        blob = bucket.blob(path)
        if not blob.exists():
            return []
        payload = blob.download_as_bytes()
        encrypted = _decode_encrypted(payload)
        plaintext = self.encryption.decrypt(encrypted)
        return json.loads(plaintext) if plaintext else []

    def _delete_blob(self, path: str) -> None:
        bucket = _get_storage_client().bucket(self.bucket_name)
        blob = bucket.blob(path)
        if blob.exists():
            blob.delete()
