"""S3 sync functionality for synchronizing pits with AWS S3."""

import json
import time
from datetime import datetime
from typing import Any, cast

import boto3
from botocore.exceptions import ClientError
from pydantic import BaseModel

from mosh.metadata import StoredObject
from mosh.pit import Pit


class S3Sync:
    """Sync pit data to/from AWS S3.

    Uploads pit data as JSON to S3 for portability across systems.
    Each object is stored as a separate JSON file under the pit prefix.
    """

    def __init__(
        self,
        bucket: str,
        prefix: str = "",
        region_name: str | None = None,
        endpoint_url: str | None = None,
    ) -> None:
        """Initialize S3 sync.

        Args:
            bucket: S3 bucket name.
            prefix: Optional prefix for S3 keys.
            region_name: AWS region name (optional).
            endpoint_url: Custom endpoint URL for S3-compatible services (optional).
        """
        self.bucket = bucket
        self.prefix = prefix.rstrip("/") if prefix else ""
        self._s3 = boto3.client(
            "s3",
            region_name=region_name,
            endpoint_url=endpoint_url,
        )

    def _get_pit_prefix(self, pit_name: str) -> str:
        """Get S3 key prefix for a pit."""
        if self.prefix:
            return f"{self.prefix}/{pit_name}"
        return pit_name

    def _get_object_key(self, pit_name: str, obj_id: str) -> str:
        """Get S3 key for an object."""
        pit_prefix = self._get_pit_prefix(pit_name)
        return f"{pit_prefix}/objects/{obj_id}.json"

    def _get_metadata_key(self, pit_name: str) -> str:
        """Get S3 key for pit metadata."""
        pit_prefix = self._get_pit_prefix(pit_name)
        return f"{pit_prefix}/metadata.json"

    def upload_object(self, pit: Pit, stored: StoredObject[BaseModel]) -> None:
        """Upload a single stored object to S3.

        Args:
            pit: The pit containing the object.
            stored: The StoredObject to upload.
        """
        key = self._get_object_key(pit.name, stored.id)
        data = {
            "id": stored.id,
            "type_name": stored.type_name,
            "created_at": stored.created_at.isoformat() if stored.created_at else None,
            "updated_at": stored.updated_at.isoformat() if stored.updated_at else None,
            "tags": stored.tags,
            "data": stored.obj.model_dump(mode="json"),
        }
        self._s3.put_object(
            Bucket=self.bucket,
            Key=key,
            Body=json.dumps(data, default=str),
            ContentType="application/json",
        )

    def download_object(self, pit_name: str, obj_id: str) -> dict[str, Any] | None:
        """Download a single object from S3.

        Args:
            pit_name: Name of the pit.
            obj_id: ID of the object to download.

        Returns:
            Object data as dict, or None if not found.
        """
        key = self._get_object_key(pit_name, obj_id)
        try:
            response = self._s3.get_object(Bucket=self.bucket, Key=key)
            return cast(dict[str, Any], json.loads(response["Body"].read().decode("utf-8")))
        except ClientError as e:
            if e.response["Error"]["Code"] == "NoSuchKey":
                return None
            raise

    def delete_object(self, pit_name: str, obj_id: str) -> None:
        """Delete a single object from S3.

        Args:
            pit_name: Name of the pit.
            obj_id: ID of the object to delete.
        """
        key = self._get_object_key(pit_name, obj_id)
        self._s3.delete_object(Bucket=self.bucket, Key=key)

    def upload_pit(self, pit: Pit) -> int:
        """Upload all objects in a pit to S3.

        Args:
            pit: The pit to upload.

        Returns:
            Number of objects uploaded.
        """
        stored_objects = pit.all_stored()
        for stored in stored_objects:
            self.upload_object(pit, stored)

        # Upload metadata
        metadata = {
            "name": pit.name,
            "object_count": len(stored_objects),
            "object_types": list(pit._object_types.keys()),
        }
        self._s3.put_object(
            Bucket=self.bucket,
            Key=self._get_metadata_key(pit.name),
            Body=json.dumps(metadata),
            ContentType="application/json",
        )

        return len(stored_objects)

    def list_objects(self, pit_name: str) -> list[str]:
        """List all object IDs in a pit on S3.

        Args:
            pit_name: Name of the pit.

        Returns:
            List of object IDs.
        """
        pit_prefix = self._get_pit_prefix(pit_name)
        objects_prefix = f"{pit_prefix}/objects/"

        object_ids = []
        paginator = self._s3.get_paginator("list_objects_v2")

        for page in paginator.paginate(Bucket=self.bucket, Prefix=objects_prefix):
            for obj in page.get("Contents", []):
                key = obj["Key"]
                # Extract object ID from key
                if key.endswith(".json"):
                    obj_id = key[len(objects_prefix) : -5]  # Remove prefix and .json
                    object_ids.append(obj_id)

        return object_ids

    def download_pit(
        self,
        pit_name: str,
        db_path: str | None = None,
        object_types: list[type[BaseModel]] | None = None,
    ) -> Pit:
        """Download a pit's data from S3.

        Args:
            pit_name: Name of the pit to download.
            db_path: Optional local database path. If None, uses in-memory database.
            object_types: List of BaseModel subclasses for deserialization.

        Returns:
            A new Pit populated with data from S3.
        """
        pit = Pit(pit_name, db_path=db_path, object_types=object_types)

        # List and download all objects
        object_ids = self.list_objects(pit_name)
        for obj_id in object_ids:
            obj_data = self.download_object(pit_name, obj_id)
            if obj_data:
                type_name = obj_data.get("type_name")
                obj_type = pit._object_types.get(type_name) if type_name else None
                if obj_type is None:
                    continue
                try:
                    obj = obj_type.model_validate(obj_data["data"])
                    tags = obj_data.get("tags", [])
                    pit.put(obj, id=obj_id, tags=tags)
                except Exception:
                    # Skip objects that can't be deserialized
                    pass

        return pit

    def sync(self, pit: Pit) -> dict[str, int]:
        """Bidirectional sync of a pit with S3.

        Uploads local objects not in S3 and downloads S3 objects not local.
        Uses object updated_at timestamps to resolve conflicts (newer wins).

        Args:
            pit: The pit to sync.

        Returns:
            Dict with counts: uploaded, downloaded, conflicts_resolved.
        """
        stats = {"uploaded": 0, "downloaded": 0, "conflicts_resolved": 0}

        # Get local and remote object IDs
        local_stored = {stored.id: stored for stored in pit.all_stored()}
        remote_ids = set(self.list_objects(pit.name))
        local_ids = set(local_stored.keys())

        # Upload objects only in local
        for obj_id in local_ids - remote_ids:
            self.upload_object(pit, local_stored[obj_id])
            stats["uploaded"] += 1

        # Download objects only in remote
        for obj_id in remote_ids - local_ids:
            obj_data = self.download_object(pit.name, obj_id)
            if obj_data:
                type_name = obj_data.get("type_name")
                obj_type = pit._object_types.get(type_name) if type_name else None
                if obj_type is None:
                    continue
                try:
                    obj = obj_type.model_validate(obj_data["data"])
                    tags = obj_data.get("tags", [])
                    pit.put(obj, id=obj_id, tags=tags)
                    stats["downloaded"] += 1
                except Exception:
                    pass

        # Handle conflicts (objects in both)
        for obj_id in local_ids & remote_ids:
            local_stored_obj = local_stored[obj_id]
            remote_data = self.download_object(pit.name, obj_id)

            if remote_data:
                remote_updated = remote_data.get("updated_at")
                if remote_updated:
                    if isinstance(remote_updated, str):
                        remote_updated = datetime.fromisoformat(remote_updated)

                local_updated = local_stored_obj.updated_at

                if remote_updated and local_updated:
                    if remote_updated > local_updated:
                        # Remote is newer, download
                        type_name = remote_data.get("type_name")
                        obj_type = pit._object_types.get(type_name) if type_name else None
                        if obj_type is None:
                            continue
                        try:
                            obj = obj_type.model_validate(remote_data["data"])
                            tags = remote_data.get("tags", [])
                            pit.put(obj, id=obj_id, tags=tags)
                            stats["downloaded"] += 1
                            stats["conflicts_resolved"] += 1
                        except Exception:
                            pass
                    elif local_updated > remote_updated:
                        # Local is newer, upload
                        self.upload_object(pit, local_stored_obj)
                        stats["uploaded"] += 1
                        stats["conflicts_resolved"] += 1

        return stats

    def sync_with_retry(
        self,
        pit: Pit,
        max_retries: int = 3,
        retry_delay: float = 5.0,
    ) -> dict[str, int]:
        """Sync with automatic retry on transient failures.

        Wraps the sync() method with retry logic for handling transient
        network failures or temporary S3 issues.

        Args:
            pit: The pit to sync.
            max_retries: Maximum number of retry attempts. Default 3.
            retry_delay: Seconds to wait between retries. Default 5.0.

        Returns:
            Dict with counts: uploaded, downloaded, conflicts_resolved.

        Raises:
            Exception: If all retry attempts fail, raises the last exception.
        """
        last_error: Exception | None = None

        for attempt in range(max_retries):
            try:
                return self.sync(pit)
            except Exception as e:
                last_error = e
                if attempt < max_retries - 1:
                    time.sleep(retry_delay)

        if last_error:
            raise last_error
        # This shouldn't happen, but satisfy type checker
        return {"uploaded": 0, "downloaded": 0, "conflicts_resolved": 0}

    def delete_pit(self, pit_name: str) -> int:
        """Delete all objects for a pit from S3.

        Args:
            pit_name: Name of the pit to delete.

        Returns:
            Number of objects deleted.
        """
        pit_prefix = self._get_pit_prefix(pit_name)
        deleted = 0

        paginator = self._s3.get_paginator("list_objects_v2")
        for page in paginator.paginate(Bucket=self.bucket, Prefix=pit_prefix):
            objects = page.get("Contents", [])
            if objects:
                delete_keys = [{"Key": obj["Key"]} for obj in objects]
                self._s3.delete_objects(
                    Bucket=self.bucket,
                    Delete={"Objects": delete_keys},
                )
                deleted += len(delete_keys)

        return deleted
