import json
import boto3
import os
from typing import Dict, List


class S3Client:
    def __init__(self):
        self.region = os.getenv("AWS_REGION", "us-east-1")
        self.client = boto3.client("s3", region_name=self.region)

    def list_buckets(self) -> List[Dict]:
        """List all S3 buckets."""
        try:
            response = self.client.list_buckets()
            return [
                {
                    "name": bucket["Name"],
                    "created": bucket["CreationDate"].isoformat()
                }
                for bucket in response.get("Buckets", [])
            ]
        except Exception as e:
            return {"error": str(e)}

    def list_objects(self, bucket: str, prefix: str = "", max_keys: int = 100) -> Dict:
        """List objects in an S3 bucket."""
        try:
            response = self.client.list_objects_v2(
                Bucket=bucket,
                Prefix=prefix,
                MaxKeys=max_keys
            )

            objects = []
            for obj in response.get("Contents", []):
                objects.append({
                    "key": obj["Key"],
                    "size": obj["Size"],
                    "modified": obj["LastModified"].isoformat(),
                    "storage_class": obj["StorageClass"]
                })

            return {
                "bucket": bucket,
                "objects": objects,
                "is_truncated": response.get("IsTruncated", False),
                "next_token": response.get("ContinuationToken")
            }
        except Exception as e:
            return {"error": str(e)}

    def get_object(self, bucket: str, key: str) -> Dict:
        """Get an object from S3."""
        try:
            response = self.client.get_object(Bucket=bucket, Key=key)
            body = response["Body"].read()

            # Try to parse as JSON if it looks like JSON
            try:
                data = json.loads(body)
                return {"key": key, "data": data, "is_json": True}
            except:
                return {
                    "key": key,
                    "data": body.decode("utf-8", errors="ignore"),
                    "is_json": False,
                    "size": len(body)
                }
        except Exception as e:
            return {"error": str(e)}

    def get_bucket_size(self, bucket: str) -> Dict:
        """Get total size of objects in a bucket."""
        try:
            paginator = self.client.get_paginator("list_objects_v2")
            pages = paginator.paginate(Bucket=bucket)

            total_size = 0
            object_count = 0

            for page in pages:
                for obj in page.get("Contents", []):
                    total_size += obj["Size"]
                    object_count += 1

            return {
                "bucket": bucket,
                "total_size_bytes": total_size,
                "object_count": object_count,
                "total_size_gb": round(total_size / (1024 ** 3), 2)
            }
        except Exception as e:
            return {"error": str(e)}