"""S3 storage backend — traces as JSON objects in S3.

Layout:
    {prefix}/traces/{project}/{trace_id}.json

Works with any S3-compatible service (AWS S3, MinIO, etc.).
"""

from __future__ import annotations

import json
import logging
import time
from typing import Any, Iterator

from ..models import Trace
from .base import StoreBase, sanitize_project_name

logger = logging.getLogger(__name__)

# Cache project list for 60 seconds to avoid repeated S3 API calls
_PROJECT_CACHE_TTL = 60


class S3Store(StoreBase):
    """S3-compatible object storage backend."""

    def __init__(
        self,
        bucket: str,
        prefix: str = "",
        endpoint_url: str | None = None,
        aws_access_key_id: str | None = None,
        aws_secret_access_key: str | None = None,
        region_name: str | None = None,
    ) -> None:
        import boto3

        self._bucket = bucket
        self._prefix = prefix.strip("/")

        kwargs: dict[str, Any] = {}
        if endpoint_url:
            kwargs["endpoint_url"] = endpoint_url
        if aws_access_key_id:
            kwargs["aws_access_key_id"] = aws_access_key_id
        if aws_secret_access_key:
            kwargs["aws_secret_access_key"] = aws_secret_access_key
        if region_name:
            kwargs["region_name"] = region_name

        self._client = boto3.client("s3", **kwargs)
        self._projects_cache: list[str] | None = None
        self._projects_cache_time: float = 0.0

    def validate(self) -> str | None:
        try:
            self._client.head_bucket(Bucket=self._bucket)
            return None
        except Exception as e:
            return f"Cannot access S3 bucket '{self._bucket}': {e}"

    # --- Key helpers ---

    def _trace_key(self, project: str, trace_id: str) -> str:
        safe_project = sanitize_project_name(project)
        parts = [self._prefix, "traces", safe_project, f"{trace_id}.json"]
        return "/".join(p for p in parts if p)

    def _traces_prefix(self, project: str | None = None) -> str:
        if project:
            safe_project = sanitize_project_name(project)
            parts = [self._prefix, "traces", safe_project, ""]
        else:
            parts = [self._prefix, "traces", ""]
        return "/".join(p for p in parts if p)

    # --- Trace operations ---

    def write_trace(self, trace: Trace) -> None:
        data = json.dumps(trace.to_dict(), indent=2)
        safe_project = sanitize_project_name(trace.project)
        self._client.put_object(
            Bucket=self._bucket,
            Key=self._trace_key(safe_project, trace.trace_id),
            Body=data.encode(),
            ContentType="application/json",
        )

    def _read_trace(self, trace_id: str) -> dict[str, Any] | None:
        # Search across project prefixes using cached project list
        for project in self._list_projects_cached():
            obj = self._get_json(self._trace_key(project, trace_id))
            if obj is not None:
                return obj
        return None

    def iter_traces(self, project: str | None = None) -> Iterator[dict[str, Any]]:
        for key in self._list_keys(self._traces_prefix(project)):
            obj = self._get_json(key)
            if obj:
                yield obj

    def _list_projects_cached(self) -> list[str]:
        """Return cached project list, refreshing if stale."""
        now = time.monotonic()
        if self._projects_cache is not None and (now - self._projects_cache_time) < _PROJECT_CACHE_TTL:
            return self._projects_cache
        projects = self.list_projects()
        self._projects_cache = projects
        self._projects_cache_time = now
        return projects

    def list_projects(self) -> list[str]:
        """List projects using S3 common prefixes (directory listing)."""
        prefix = self._traces_prefix()
        projects = []
        paginator = self._client.get_paginator("list_objects_v2")
        for page in paginator.paginate(
            Bucket=self._bucket, Prefix=prefix, Delimiter="/"
        ):
            for cp in page.get("CommonPrefixes", []):
                # cp["Prefix"] looks like "{prefix}/traces/{project}/"
                project = cp["Prefix"].rstrip("/").rsplit("/", 1)[-1]
                if project:
                    projects.append(project)
        result = sorted(projects)
        self._projects_cache = result
        self._projects_cache_time = time.monotonic()
        return result

    # --- S3 helpers ---

    def _get_json(self, key: str) -> dict[str, Any] | None:
        try:
            response = self._client.get_object(Bucket=self._bucket, Key=key)
            return json.loads(response["Body"].read())
        except self._client.exceptions.NoSuchKey:
            return None
        except Exception as e:
            logger.warning(f"S3 read error for {key}: {e}")
            return None

    def _list_keys(self, prefix: str) -> list[str]:
        """List all object keys under a prefix."""
        keys: list[str] = []
        paginator = self._client.get_paginator("list_objects_v2")
        for page in paginator.paginate(Bucket=self._bucket, Prefix=prefix):
            for obj in page.get("Contents", []):
                if obj["Key"].endswith(".json"):
                    keys.append(obj["Key"])
        return keys
