"""
FetchDocumentOperator - Fetch documents from SeaweedFS/S3 storage.

This operator downloads documents from S3-compatible storage (SeaweedFS)
to local temporary storage for processing.
"""

from airflow.models import BaseOperator
import tempfile
import os
import shutil
from typing import Dict

from api.s3 import download_file


class FetchDocumentOperator(BaseOperator):
    """
    Fetch document from SeaweedFS/S3 storage.

    Input: ValidatedEvent from validate_event task (via XCom)
    Output: FetchedDocument dict via XCom

    Uses boto3 client for S3-compatible SeaweedFS storage access.
    Credentials are read from AWS environment variables:
    - AWS_ACCESS_KEY_ID
    - AWS_SECRET_ACCESS_KEY
    - AWS_ENDPOINT_URL
    """

    def __init__(
        self,
        validate_event_task_id: str = 'validate_event',
        **kwargs
    ):
        """
        Initialize FetchDocumentOperator.

        Args:
            validate_event_task_id: Task ID of the validation task to pull XCom from
            **kwargs: Additional BaseOperator arguments
        """
        super().__init__(**kwargs)
        self.validate_event_task_id = validate_event_task_id

    def execute(self, context) -> Dict:
        """
        Download document from storage to local temp directory.

        Args:
            context: Airflow task context with TaskInstance

        Returns:
            Dict: FetchedDocument schema containing:
                - local_path: Local filesystem path to downloaded file
                - content_type: MIME type of the document
                - original_filename: Original filename from storage
                - file_extension: File extension (e.g., 'xlsx', 'pdf')
                - document_id: Document identifier

        Raises:
            ValueError: If event data is invalid or download fails
            KeyError: If required fields missing from event data
        """
        # Get event from previous task
        event = context['ti'].xcom_pull(task_ids=self.validate_event_task_id)

        if not event:
            raise ValueError(
                f"No event data found from task: {self.validate_event_task_id}"
            )

        # Validate required fields (Issue #4 fix)
        required_fields = ['bucket', 'object_key', 'document_id', 'content_type']
        missing = [f for f in required_fields if f not in event]
        if missing:
            raise ValueError(f"Missing required fields in event: {', '.join(missing)}")

        bucket = event['bucket']
        object_key = event['object_key']
        document_id = event['document_id']
        tenant_id = event.get('tenant_id', '')
        node_id = event.get('node_id', '')

        self.log.info(
            f"Fetching document {document_id} from bucket '{bucket}', "
            f"key: {object_key}, tenant: {tenant_id}, node: {node_id}"
        )

        # Validate tenant/node prefix if provided (data-isolation-spec Section 4.4)
        if tenant_id and node_id:
            expected_prefix = f"{tenant_id}/{node_id}/"
            if not object_key.startswith(expected_prefix):
                raise ValueError(
                    f"Object key '{object_key}' does not match tenant/node scope "
                    f"'{expected_prefix}'. Cross-tenant file access blocked."
                )

        # Extract filename and extension
        original_filename = os.path.basename(object_key)
        file_extension = original_filename.lower().split('.')[-1] if '.' in original_filename else ''

        # Create temp directory for download
        # Note: Temp files are intentionally NOT cleaned up here
        # Airflow tasks may need the files for downstream tasks
        # Cleanup should be handled by a dedicated cleanup task or Airflow's tmpdir cleanup
        temp_dir = tempfile.mkdtemp(prefix=f'airflow_doc_{document_id}_')
        local_path = os.path.join(temp_dir, original_filename)

        try:
            self.log.info(f"Downloading to local path: {local_path}")

            # Download via the breaker-wrapped S3 helper (api.s3)
            download_file(bucket, object_key, local_path)

            file_size = os.path.getsize(local_path)
            self.log.info(
                f"Successfully downloaded document to: {local_path} "
                f"({file_size} bytes)"
            )

        except Exception as e:
            self.log.error(f"Failed to download document: {str(e)}")
            # Clean up temp directory on failure (Issue #2 fix)
            if os.path.exists(temp_dir):
                shutil.rmtree(temp_dir, ignore_errors=True)
            raise

        return {
            'local_path': local_path,
            'content_type': event['content_type'],
            'original_filename': original_filename,
            'file_extension': file_extension,
            'document_id': document_id,
            'tenant_id': tenant_id,
            'node_id': node_id,
        }