"""S3 service emulator."""

import io
import logging
import re
from collections import defaultdict
from dataclasses import dataclass, field
from typing import Optional
from urllib.parse import unquote

from flask import Request, Response

from localrun.utils import (
    error_response,
    etag,
    iso_timestamp,
    md5_hex,
    new_request_id,
    xml_response,
)

logger = logging.getLogger("localrun.s3")


@dataclass
class S3Object:
    """An object stored in S3."""
    key: str
    data: bytes
    content_type: str = "application/octet-stream"
    metadata: dict = field(default_factory=dict)
    etag: str = ""
    last_modified: str = ""
    size: int = 0

    def __post_init__(self):
        if not self.etag:
            self.etag = f'"{md5_hex(self.data)}"'
        if not self.last_modified:
            self.last_modified = iso_timestamp()
        self.size = len(self.data)


class S3Service:
    """In-memory S3 emulator."""

    def __init__(self):
        # buckets[bucket_name] = {key: S3Object}
        self.buckets: dict[str, dict[str, S3Object]] = {}
        self.bucket_created: dict[str, str] = {}  # bucket_name -> creation date

    def handle(self, req: Request, path: str) -> Response:
        """Route S3 requests."""
        parts = path.strip("/").split("/", 1) if path.strip("/") else []
        bucket_name = parts[0] if parts else None
        object_key = unquote(parts[1]) if len(parts) > 1 else None

        method = req.method

        # List all buckets
        if not bucket_name:
            if method == "GET":
                return self._list_buckets(req)
            return error_response("MethodNotAllowed", "Method not allowed", 405)

        # Bucket-level operations (no key)
        if object_key is None:
            # Check for ?delete (multi-delete)
            if "delete" in req.args:
                return self._delete_objects(req, bucket_name)
            # Check for ?location
            if "location" in req.args:
                return self._get_bucket_location(req, bucket_name)
            if method == "PUT":
                return self._create_bucket(req, bucket_name)
            if method == "DELETE":
                return self._delete_bucket(req, bucket_name)
            if method == "GET":
                return self._list_objects(req, bucket_name)
            if method == "HEAD":
                return self._head_bucket(req, bucket_name)

        # Object-level operations
        if object_key is not None:
            if method == "PUT":
                # Check for x-amz-copy-source (CopyObject)
                copy_source = req.headers.get("x-amz-copy-source")
                if copy_source:
                    return self._copy_object(req, bucket_name, object_key, copy_source)
                return self._put_object(req, bucket_name, object_key)
            if method == "GET":
                return self._get_object(req, bucket_name, object_key)
            if method == "DELETE":
                return self._delete_object(req, bucket_name, object_key)
            if method == "HEAD":
                return self._head_object(req, bucket_name, object_key)

        return error_response("NotImplemented", "Operation not implemented", 501)

    # ─── Bucket operations ────────────────────────────────────────────────────

    def _list_buckets(self, req: Request) -> Response:
        buckets_xml = ""
        for name in sorted(self.buckets.keys()):
            created = self.bucket_created.get(name, iso_timestamp())
            buckets_xml += f"""
        <Bucket>
            <Name>{name}</Name>
            <CreationDate>{created}</CreationDate>
        </Bucket>"""

        body = f"""<?xml version="1.0" encoding="UTF-8"?>
<ListAllMyBucketsResult xmlns="http://s3.amazonaws.com/doc/2006-03-01/">
    <Owner>
        <ID>000000000000</ID>
        <DisplayName>localrun</DisplayName>
    </Owner>
    <Buckets>{buckets_xml}
    </Buckets>
</ListAllMyBucketsResult>"""
        return Response(body, status=200, content_type="application/xml")

    def _create_bucket(self, req: Request, bucket_name: str) -> Response:
        if bucket_name in self.buckets:
            return error_response(
                "BucketAlreadyOwnedByYou",
                "Your previous request to create the named bucket succeeded.",
                status=409,
            )
        self.buckets[bucket_name] = {}
        self.bucket_created[bucket_name] = iso_timestamp()
        logger.info("Created bucket: %s", bucket_name)
        resp = Response("", status=200)
        resp.headers["Location"] = f"/{bucket_name}"
        return resp

    def _delete_bucket(self, req: Request, bucket_name: str) -> Response:
        if bucket_name not in self.buckets:
            return error_response("NoSuchBucket", f"Bucket '{bucket_name}' does not exist.", 404)
        if self.buckets[bucket_name]:
            return error_response("BucketNotEmpty", "The bucket you tried to delete is not empty.", 409)
        del self.buckets[bucket_name]
        self.bucket_created.pop(bucket_name, None)
        logger.info("Deleted bucket: %s", bucket_name)
        return Response("", status=204)

    def _head_bucket(self, req: Request, bucket_name: str) -> Response:
        if bucket_name not in self.buckets:
            return error_response("NoSuchBucket", f"Bucket '{bucket_name}' does not exist.", 404)
        return Response("", status=200)

    def _get_bucket_location(self, req: Request, bucket_name: str) -> Response:
        if bucket_name not in self.buckets:
            return error_response("NoSuchBucket", f"Bucket '{bucket_name}' does not exist.", 404)
        body = """<?xml version="1.0" encoding="UTF-8"?>
<LocationConstraint xmlns="http://s3.amazonaws.com/doc/2006-03-01/">us-east-1</LocationConstraint>"""
        return Response(body, status=200, content_type="application/xml")

    # ─── Object operations ────────────────────────────────────────────────────

    def _list_objects(self, req: Request, bucket_name: str) -> Response:
        if bucket_name not in self.buckets:
            return error_response("NoSuchBucket", f"Bucket '{bucket_name}' does not exist.", 404)

        prefix = req.args.get("prefix", "")
        delimiter = req.args.get("delimiter", "")
        max_keys = int(req.args.get("max-keys", "1000"))
        list_type = req.args.get("list-type", "1")

        objects = self.buckets[bucket_name]
        matching = []
        common_prefixes = set()

        for key, obj in sorted(objects.items()):
            if prefix and not key.startswith(prefix):
                continue
            if delimiter:
                rest = key[len(prefix):]
                idx = rest.find(delimiter)
                if idx >= 0:
                    common_prefixes.add(prefix + rest[: idx + len(delimiter)])
                    continue
            matching.append(obj)

        matching = matching[:max_keys]

        contents_xml = ""
        for obj in matching:
            contents_xml += f"""
        <Contents>
            <Key>{obj.key}</Key>
            <LastModified>{obj.last_modified}</LastModified>
            <ETag>{obj.etag}</ETag>
            <Size>{obj.size}</Size>
            <StorageClass>STANDARD</StorageClass>
        </Contents>"""

        prefixes_xml = ""
        for cp in sorted(common_prefixes):
            prefixes_xml += f"""
        <CommonPrefixes>
            <Prefix>{cp}</Prefix>
        </CommonPrefixes>"""

        if list_type == "2":
            body = f"""<?xml version="1.0" encoding="UTF-8"?>
<ListBucketResult xmlns="http://s3.amazonaws.com/doc/2006-03-01/">
    <Name>{bucket_name}</Name>
    <Prefix>{prefix}</Prefix>
    <KeyCount>{len(matching)}</KeyCount>
    <MaxKeys>{max_keys}</MaxKeys>
    <IsTruncated>false</IsTruncated>{contents_xml}{prefixes_xml}
</ListBucketResult>"""
        else:
            body = f"""<?xml version="1.0" encoding="UTF-8"?>
<ListBucketResult xmlns="http://s3.amazonaws.com/doc/2006-03-01/">
    <Name>{bucket_name}</Name>
    <Prefix>{prefix}</Prefix>
    <MaxKeys>{max_keys}</MaxKeys>
    <IsTruncated>false</IsTruncated>{contents_xml}{prefixes_xml}
</ListBucketResult>"""

        return Response(body, status=200, content_type="application/xml")

    def _put_object(self, req: Request, bucket_name: str, key: str) -> Response:
        if bucket_name not in self.buckets:
            return error_response("NoSuchBucket", f"Bucket '{bucket_name}' does not exist.", 404)

        data = req.get_data()
        content_type = req.headers.get("Content-Type", "application/octet-stream")

        # Extract user metadata (x-amz-meta-*)
        metadata = {}
        for header, value in req.headers:
            if header.lower().startswith("x-amz-meta-"):
                meta_key = header[len("x-amz-meta-"):]
                metadata[meta_key] = value

        obj = S3Object(
            key=key,
            data=data,
            content_type=content_type,
            metadata=metadata,
        )
        self.buckets[bucket_name][key] = obj
        logger.info("Put object: s3://%s/%s (%d bytes)", bucket_name, key, len(data))

        resp = Response("", status=200)
        resp.headers["ETag"] = obj.etag
        return resp

    def _get_object(self, req: Request, bucket_name: str, key: str) -> Response:
        if bucket_name not in self.buckets:
            return error_response("NoSuchBucket", f"Bucket '{bucket_name}' does not exist.", 404)
        obj = self.buckets[bucket_name].get(key)
        if not obj:
            return error_response("NoSuchKey", f"The specified key does not exist: {key}", 404)

        resp = Response(obj.data, status=200, content_type=obj.content_type)
        resp.headers["ETag"] = obj.etag
        resp.headers["Content-Length"] = str(obj.size)
        resp.headers["Last-Modified"] = obj.last_modified
        for mk, mv in obj.metadata.items():
            resp.headers[f"x-amz-meta-{mk}"] = mv
        return resp

    def _head_object(self, req: Request, bucket_name: str, key: str) -> Response:
        if bucket_name not in self.buckets:
            return error_response("NoSuchBucket", f"Bucket '{bucket_name}' does not exist.", 404)
        obj = self.buckets[bucket_name].get(key)
        if not obj:
            return error_response("NoSuchKey", f"The specified key does not exist: {key}", 404)

        resp = Response("", status=200, content_type=obj.content_type)
        resp.headers["ETag"] = obj.etag
        resp.headers["Content-Length"] = str(obj.size)
        resp.headers["Last-Modified"] = obj.last_modified
        return resp

    def _delete_object(self, req: Request, bucket_name: str, key: str) -> Response:
        if bucket_name not in self.buckets:
            return error_response("NoSuchBucket", f"Bucket '{bucket_name}' does not exist.", 404)
        self.buckets[bucket_name].pop(key, None)
        logger.info("Deleted object: s3://%s/%s", bucket_name, key)
        return Response("", status=204)

    def _copy_object(self, req: Request, dest_bucket: str, dest_key: str, copy_source: str) -> Response:
        # Parse source: /bucket/key
        source = copy_source.lstrip("/")
        parts = source.split("/", 1)
        if len(parts) < 2:
            return error_response("InvalidArgument", "Invalid copy source", 400)
        src_bucket, src_key = parts[0], unquote(parts[1])

        if src_bucket not in self.buckets:
            return error_response("NoSuchBucket", f"Source bucket '{src_bucket}' does not exist.", 404)
        src_obj = self.buckets[src_bucket].get(src_key)
        if not src_obj:
            return error_response("NoSuchKey", f"Source key does not exist: {src_key}", 404)
        if dest_bucket not in self.buckets:
            return error_response("NoSuchBucket", f"Destination bucket '{dest_bucket}' does not exist.", 404)

        new_obj = S3Object(
            key=dest_key,
            data=src_obj.data,
            content_type=src_obj.content_type,
            metadata=dict(src_obj.metadata),
        )
        self.buckets[dest_bucket][dest_key] = new_obj

        body = f"""<?xml version="1.0" encoding="UTF-8"?>
<CopyObjectResult>
    <ETag>{new_obj.etag}</ETag>
    <LastModified>{new_obj.last_modified}</LastModified>
</CopyObjectResult>"""
        return Response(body, status=200, content_type="application/xml")

    def _delete_objects(self, req: Request, bucket_name: str) -> Response:
        """Handle multi-object delete."""
        if bucket_name not in self.buckets:
            return error_response("NoSuchBucket", f"Bucket '{bucket_name}' does not exist.", 404)

        import xml.etree.ElementTree as ET
        try:
            root = ET.fromstring(req.get_data())
        except ET.ParseError:
            return error_response("MalformedXML", "The XML you provided was not well-formed.", 400)

        ns = "http://s3.amazonaws.com/doc/2006-03-01/"
        deleted_xml = ""
        for obj_elem in root.findall(f".//{{{ns}}}Object") or root.findall(".//Object") or root:
            key_elem = obj_elem.find(f"{{{ns}}}Key")
            if key_elem is None:
                key_elem = obj_elem.find("Key")
            if key_elem is not None and key_elem.text:
                key = key_elem.text
                self.buckets[bucket_name].pop(key, None)
                deleted_xml += f"""
        <Deleted>
            <Key>{key}</Key>
        </Deleted>"""

        body = f"""<?xml version="1.0" encoding="UTF-8"?>
<DeleteResult xmlns="http://s3.amazonaws.com/doc/2006-03-01/">{deleted_xml}
</DeleteResult>"""
        return Response(body, status=200, content_type="application/xml")
