"""SeaweedFS Image Upload Service.

This service handles uploading extracted images to SeaweedFS/S3-compatible storage
and returns URLs/keys for reference in Milvus metadata.

Key format convention (data-isolation-spec Section 4.3):
  {tenant_id}/{node_id}/{document_id}/images/{sheet_name}/{image_number}.{ext}
"""

import os
import boto3
import base64
import logging
from typing import Dict, List, Optional, Any
from uuid import UUID
from io import BytesIO

logger = logging.getLogger(__name__)

DEFAULT_IMAGE_BUCKET = os.environ.get("IMAGE_BUCKET", "textiq-images")


def validate_s3_path(tenant_id: str, node_id: str, s3_key: str) -> None:
    """Reject any S3 operation where the key does not match tenant/node scope.

    Data isolation spec Section 4.5.

    Args:
        tenant_id: Expected tenant UUID string.
        node_id: Expected node UUID string.
        s3_key: The S3 object key to validate.

    Raises:
        PermissionError: If key does not start with expected tenant/node prefix.
    """
    expected_prefix = f"{tenant_id}/{node_id}/"
    if not s3_key.startswith(expected_prefix):
        raise PermissionError(
            f"S3 key '{s3_key}' does not match scope '{expected_prefix}'. "
            "Cross-tenant/cross-node file access blocked."
        )


class SeaweedFSService:
    """Service for uploading images to SeaweedFS/S3-compatible storage."""

    def __init__(self):
        """Initialize SeaweedFS service with boto3 S3 client.

        Uses environment variables for configuration:
        - AWS_ACCESS_KEY_ID
        - AWS_SECRET_ACCESS_KEY
        - AWS_ENDPOINT_URL
        """
        # Configure boto3 for SeaweedFS with signature version v4
        # This is required for S3-compatible storage
        from botocore.client import Config

        self.s3_client = boto3.client(
            's3',
            config=Config(signature_version='s3v4'),
            region_name='us-east-1',
        )
        logger.info("Initialized SeaweedFS service")

    def _ensure_bucket_exists(self, bucket: str) -> bool:
        """Ensure the bucket exists, create if not.

        Args:
            bucket: Bucket name

        Returns:
            True if bucket exists or was created, False otherwise
        """
        try:
            self.s3_client.head_bucket(Bucket=bucket)
            return True
        except Exception as e:
            # Bucket doesn't exist, try to create it
            try:
                self.s3_client.create_bucket(Bucket=bucket)
                logger.info(f"Created bucket: {bucket}")
                return True
            except Exception as create_err:
                logger.error(f"Failed to create bucket {bucket}: {create_err}")
                return False

    def upload_image(
        self,
        image_base64: str,
        document_id: UUID,
        image_number: int,
        sheet_name: str,
        mime_type: str = "image/png",
        bucket: str = DEFAULT_IMAGE_BUCKET,
        tenant_id: str = "",
        node_id: str = "",
    ) -> Dict[str, str]:
        """Upload a single image to SeaweedFS.

        Args:
            image_base64: Base64-encoded image data
            document_id: Document UUID
            image_number: Sequential image number within document
            sheet_name: Sheet name where image was found
            mime_type: MIME type of image (e.g., "image/png")
            bucket: S3 bucket name
            tenant_id: Tenant UUID string for key prefix (data isolation)
            node_id: Node UUID string for key prefix (data isolation)

        Returns:
            Dict with upload result:
                - image_s3_key: S3 key/path of uploaded image
                - image_url: Full URL to access the image
                - success: Boolean indicating upload success

        Raises:
            Exception: If upload fails
        """
        # Ensure bucket exists before uploading
        if not self._ensure_bucket_exists(bucket):
            return {
                "image_s3_key": "",
                "image_url": "",
                "success": False,
                "error": f"Bucket {bucket} does not exist and could not be created",
            }

        try:
            # Decode base64 to bytes
            image_bytes = base64.b64decode(image_base64)

            # Determine file extension from MIME type
            ext_map = {
                "image/png": "png",
                "image/jpeg": "jpg",
                "image/jpg": "jpg",
                "image/gif": "gif",
                "image/bmp": "bmp",
                "image/tiff": "tiff",
            }
            extension = ext_map.get(mime_type, "png")

            # Generate S3 key with tenant-scoped hierarchical structure
            # Pattern: {tenant_id}/{node_id}/{doc_id}/images/{sheet_name}/{image_number}.{ext}
            safe_sheet_name = "".join(c if c.isalnum() or c in ('-', '_') else '_' for c in sheet_name)
            if tenant_id and node_id:
                s3_key = f"{tenant_id}/{node_id}/{document_id}/images/{safe_sheet_name}/image_{image_number}.{extension}"
                validate_s3_path(tenant_id, node_id, s3_key)
            else:
                # Legacy fallback for callers that don't provide tenant context yet
                s3_key = f"documents/{document_id}/images/{safe_sheet_name}/image_{image_number}.{extension}"

            # Upload to SeaweedFS via S3 API
            self.s3_client.put_object(
                Bucket=bucket,
                Key=s3_key,
                Body=BytesIO(image_bytes),
                ContentType=mime_type,
            )

            # Construct URL (SeaweedFS URL pattern)
            endpoint_url = self.s3_client.meta.endpoint_url
            image_url = f"{endpoint_url}/{bucket}/{s3_key}"

            logger.info(f"Successfully uploaded image to {s3_key}")

            return {
                "image_s3_key": s3_key,
                "image_url": image_url,
                "success": True,
            }

        except Exception as e:
            logger.error(f"Failed to upload image {image_number} for document {document_id}: {e}")
            return {
                "image_s3_key": "",
                "image_url": "",
                "success": False,
                "error": str(e),
            }

    def upload_images_batch(
        self,
        chunks: List[Dict[str, Any]],
        document_id: UUID,
        bucket: str = DEFAULT_IMAGE_BUCKET,
        tenant_id: str = "",
        node_id: str = "",
    ) -> List[Dict[str, Any]]:
        """Upload all images from DETAILED pipeline chunks to SeaweedFS.

        Processes chunks, uploads images with base64 data, and updates
        metadata with SeaweedFS URLs/keys.

        Args:
            chunks: List of chunk dicts (may contain image_base64 in metadata)
            document_id: Document UUID
            bucket: S3 bucket name
            tenant_id: Tenant UUID string for key prefix (data isolation)
            node_id: Node UUID string for key prefix (data isolation)

        Returns:
            Updated chunks with image_url and image_s3_key added,
            and image_base64 removed after successful upload
        """
        updated_chunks = []
        image_counter = 0

        for chunk in chunks:
            # Get metadata from _legacy_metadata if available
            legacy_meta = chunk.get('_legacy_metadata', {})
            full_metadata = legacy_meta.get('full_metadata', {})
            element_type = legacy_meta.get('element_type', 'Text')

            # Check if this chunk contains an image
            if element_type == 'Image' and 'image_base64' in full_metadata:
                image_counter += 1
                image_base64 = full_metadata['image_base64']
                sheet_name = chunk.get('_source', {}).get('sheet', 'Sheet1')
                mime_type = full_metadata.get('mime_type', 'image/png')

                logger.info(f"Uploading image {image_counter} from sheet '{sheet_name}'")

                # Upload to SeaweedFS
                upload_result = self.upload_image(
                    image_base64=image_base64,
                    document_id=document_id,
                    image_number=image_counter,
                    sheet_name=sheet_name,
                    mime_type=mime_type,
                    bucket=bucket,
                    tenant_id=tenant_id,
                    node_id=node_id,
                )

                if upload_result['success']:
                    # Update metadata with SeaweedFS info
                    full_metadata['image_s3_key'] = upload_result['image_s3_key']
                    full_metadata['image_url'] = upload_result['image_url']

                    # Remove base64 data after successful upload
                    full_metadata.pop('image_base64', None)

                    logger.info(f"✓ Image {image_counter} uploaded: {upload_result['image_s3_key']}")
                else:
                    logger.warning(f"✗ Failed to upload image {image_counter}: {upload_result.get('error', 'Unknown error')}")

            updated_chunks.append(chunk)

        logger.info(f"Processed {image_counter} images from {len(chunks)} chunks")
        return updated_chunks

    def delete_document_images(
        self,
        document_id: UUID,
        bucket: str = DEFAULT_IMAGE_BUCKET,
        tenant_id: str = "",
        node_id: str = "",
    ) -> Dict[str, Any]:
        """Delete all images for a document from SeaweedFS.

        Args:
            document_id: Document UUID
            bucket: S3 bucket name
            tenant_id: Tenant UUID string for scoped deletion
            node_id: Node UUID string for scoped deletion

        Returns:
            Dict with deletion result:
                - deleted_count: Number of images deleted
                - success: Boolean indicating success
        """
        try:
            # List all objects with the tenant-scoped document prefix
            if tenant_id and node_id:
                prefix = f"{tenant_id}/{node_id}/{document_id}/images/"
            else:
                prefix = f"documents/{document_id}/images/"
            response = self.s3_client.list_objects_v2(
                Bucket=bucket,
                Prefix=prefix
            )

            if 'Contents' not in response:
                logger.info(f"No images found for document {document_id}")
                return {"deleted_count": 0, "success": True}

            # Delete all objects
            objects_to_delete = [{'Key': obj['Key']} for obj in response['Contents']]

            if objects_to_delete:
                self.s3_client.delete_objects(
                    Bucket=bucket,
                    Delete={'Objects': objects_to_delete}
                )

                deleted_count = len(objects_to_delete)
                logger.info(f"Deleted {deleted_count} images for document {document_id}")

                return {
                    "deleted_count": deleted_count,
                    "success": True,
                }

            return {"deleted_count": 0, "success": True}

        except Exception as e:
            logger.error(f"Failed to delete images for document {document_id}: {e}")
            return {
                "deleted_count": 0,
                "success": False,
                "error": str(e),
            }