"""
Copyright 2024-2026 ChatterMate

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import base64
import os
import uuid
from typing import Dict, Optional, List, Set
from app.core.config import settings
from app.core.s3 import upload_file_to_s3
from app.services.file_storage import resolve_public_url
from app.core.logger import get_logger

logger = get_logger(__name__)

# Allowed file types and size limits
ALLOWED_IMAGE_TYPES = {'image/jpeg', 'image/jpg', 'image/png', 'image/gif', 'image/webp'}
ALLOWED_DOCUMENT_TYPES = {'application/pdf', 'application/msword', 
                          'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
                          'text/plain', 'text/csv', 'application/vnd.ms-excel',
                          'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'}
ALLOWED_FILE_TYPES = ALLOWED_IMAGE_TYPES | ALLOWED_DOCUMENT_TYPES

# Security: SVG removed from default allowed types due to XSS risks
# Can be enabled per-agent if needed with proper sanitization

MAX_FILE_SIZE = 10 * 1024 * 1024  # 10MB
MAX_IMAGE_SIZE = 5 * 1024 * 1024  # 5MB

# Magic byte signatures for file type validation
# This prevents MIME type spoofing attacks
MAGIC_BYTES = {
    # Images
    'image/jpeg': [
        (b'\xff\xd8\xff\xe0', 0),  # JPEG JFIF
        (b'\xff\xd8\xff\xe1', 0),  # JPEG EXIF
        (b'\xff\xd8\xff\xe2', 0),  # JPEG ICC
        (b'\xff\xd8\xff\xe8', 0),  # JPEG SPIFF
        (b'\xff\xd8\xff\xdb', 0),  # JPEG DQT
        (b'\xff\xd8\xff\xee', 0),  # JPEG Adobe
    ],
    'image/jpg': [
        (b'\xff\xd8\xff\xe0', 0),
        (b'\xff\xd8\xff\xe1', 0),
        (b'\xff\xd8\xff\xe2', 0),
        (b'\xff\xd8\xff\xe8', 0),
        (b'\xff\xd8\xff\xdb', 0),
        (b'\xff\xd8\xff\xee', 0),
    ],
    'image/png': [
        (b'\x89PNG\r\n\x1a\n', 0),
    ],
    'image/gif': [
        (b'GIF87a', 0),
        (b'GIF89a', 0),
    ],
    'image/webp': [
        (b'RIFF', 0),  # WebP starts with RIFF, then has WEBP at offset 8
    ],
    # Documents
    'application/pdf': [
        (b'%PDF-', 0),
    ],
    'application/msword': [
        (b'\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1', 0),  # OLE Compound File (DOC)
    ],
    'application/vnd.openxmlformats-officedocument.wordprocessingml.document': [
        (b'PK\x03\x04', 0),  # ZIP-based (DOCX)
    ],
    'application/vnd.ms-excel': [
        (b'\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1', 0),  # OLE Compound File (XLS)
    ],
    'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet': [
        (b'PK\x03\x04', 0),  # ZIP-based (XLSX)
    ],
    'text/plain': [],  # Text files don't have magic bytes, validated by content
    'text/csv': [],  # CSV files don't have magic bytes, validated by content
}

# File extension to MIME type mapping for validation
EXTENSION_TO_MIME = {
    '.jpg': 'image/jpeg',
    '.jpeg': 'image/jpeg',
    '.png': 'image/png',
    '.gif': 'image/gif',
    '.webp': 'image/webp',
    '.pdf': 'application/pdf',
    '.doc': 'application/msword',
    '.docx': 'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
    '.txt': 'text/plain',
    '.csv': 'text/csv',
    '.xls': 'application/vnd.ms-excel',
    '.xlsx': 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
}

# Category labels for display
FILE_TYPE_CATEGORIES = {
    'images': {'image/jpeg', 'image/jpg', 'image/png', 'image/gif', 'image/webp'},
    'documents': {'application/pdf'},
    'office': {'application/msword', 'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
               'application/vnd.ms-excel', 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'},
    'text': {'text/plain', 'text/csv'},
}


class FileUploadService:
    """Service for handling file uploads to S3 or local storage with enhanced security"""
    
    @staticmethod
    def validate_magic_bytes(file_content: bytes, claimed_content_type: str) -> tuple[bool, Optional[str]]:
        """
        Validate file content against magic byte signatures to prevent MIME type spoofing.
        
        Args:
            file_content: The raw binary content of the file
            claimed_content_type: The MIME type claimed by the client
            
        Returns:
            (is_valid, error_message)
        """
        if claimed_content_type not in MAGIC_BYTES:
            logger.warning(f"Unknown content type for magic byte validation: {claimed_content_type}")
            return False, f"File type '{claimed_content_type}' is not allowed"
        
        signatures = MAGIC_BYTES[claimed_content_type]
        
        # Text files don't have magic bytes - validate as text content
        if not signatures:
            if claimed_content_type in {'text/plain', 'text/csv'}:
                # Check if content is valid text (not binary)
                try:
                    # Try to decode as UTF-8, allowing for some common encodings
                    sample = file_content[:8192]  # Check first 8KB
                    sample.decode('utf-8')
                    return True, None
                except UnicodeDecodeError:
                    try:
                        sample.decode('latin-1')
                        return True, None
                    except UnicodeDecodeError:
                        return False, "File content does not appear to be valid text"
            return True, None
        
        # Check for WebP special case (has WEBP at offset 8)
        if claimed_content_type == 'image/webp':
            if len(file_content) >= 12:
                if file_content[:4] == b'RIFF' and file_content[8:12] == b'WEBP':
                    return True, None
            return False, "File content does not match WebP format signature"
        
        # Check against all possible signatures for this MIME type
        for signature, offset in signatures:
            end_offset = offset + len(signature)
            if len(file_content) >= end_offset:
                if file_content[offset:end_offset] == signature:
                    return True, None
        
        return False, f"File content does not match the claimed type '{claimed_content_type}'. The file may be corrupted or mislabeled."
    
    @staticmethod
    def validate_extension(filename: str, claimed_content_type: str) -> tuple[bool, Optional[str]]:
        """
        Validate that file extension matches the claimed content type.
        
        Args:
            filename: The name of the file
            claimed_content_type: The MIME type claimed by the client
            
        Returns:
            (is_valid, error_message)
        """
        ext = os.path.splitext(filename.lower())[1]
        
        if not ext:
            return False, "File must have an extension"
        
        if ext not in EXTENSION_TO_MIME:
            return False, f"File extension '{ext}' is not allowed. Allowed extensions: {', '.join(sorted(EXTENSION_TO_MIME.keys()))}"
        
        expected_mime = EXTENSION_TO_MIME[ext]
        
        # Handle jpg/jpeg equivalence
        if claimed_content_type in {'image/jpeg', 'image/jpg'} and expected_mime in {'image/jpeg', 'image/jpg'}:
            return True, None
        
        if expected_mime != claimed_content_type:
            return False, f"File extension '{ext}' does not match content type '{claimed_content_type}'"
        
        return True, None
    
    @staticmethod
    def validate_file_type(content_type: str, allowed_types: Optional[Set[str]] = None) -> bool:
        """
        Validate if the file type is allowed.
        
        Args:
            content_type: The MIME type to validate
            allowed_types: Optional set of allowed MIME types. If None, uses default ALLOWED_FILE_TYPES.
        """
        types_to_check = allowed_types if allowed_types is not None else ALLOWED_FILE_TYPES
        return content_type in types_to_check
    
    @staticmethod
    def get_allowed_types_for_agent(agent_allowed_types: Optional[List[str]] = None) -> Set[str]:
        """
        Get the set of allowed file types for an agent.
        
        Args:
            agent_allowed_types: List of category names like ['images', 'documents', 'office', 'text']
                                 If None or empty, returns all default allowed types.
        
        Returns:
            Set of MIME types that are allowed
        """
        if not agent_allowed_types:
            return ALLOWED_FILE_TYPES
        
        allowed = set()
        for category in agent_allowed_types:
            category_lower = category.lower()
            if category_lower in FILE_TYPE_CATEGORIES:
                allowed.update(FILE_TYPE_CATEGORIES[category_lower])
        
        # If no valid categories found, return default
        if not allowed:
            return ALLOWED_FILE_TYPES
        
        return allowed
    
    @staticmethod
    def validate_file_size(file_size: int, content_type: str) -> tuple[bool, Optional[str]]:
        """
        Validate file size based on content type
        Returns: (is_valid, error_message)
        """
        if content_type in ALLOWED_IMAGE_TYPES and file_size > MAX_IMAGE_SIZE:
            return False, f"Image file size exceeds maximum allowed size of {MAX_IMAGE_SIZE / (1024*1024):.1f}MB"
        
        if file_size > MAX_FILE_SIZE:
            return False, f"File size exceeds maximum allowed size of {MAX_FILE_SIZE / (1024*1024):.1f}MB"
        
        return True, None
    
    @staticmethod
    def get_friendly_allowed_types_message(allowed_types: Optional[Set[str]] = None) -> str:
        """Generate a user-friendly message listing allowed file types."""
        types_to_use = allowed_types if allowed_types is not None else ALLOWED_FILE_TYPES
        
        extensions = []
        for ext, mime in EXTENSION_TO_MIME.items():
            if mime in types_to_use:
                extensions.append(ext.upper().replace('.', ''))
        
        return ', '.join(sorted(set(extensions)))
    
    @staticmethod
    async def upload_file(
        file_data: dict,
        org_id: str,
        customer_id: Optional[str] = None,
        user_id: Optional[str] = None,
        allowed_types: Optional[Set[str]] = None
    ) -> Dict[str, any]:
        """
        Upload a file to S3 or local storage with enhanced security validation.
        
        Args:
            file_data: Dict with 'content' (base64), 'filename', 'content_type', 'size'
            org_id: Organization ID
            customer_id: Optional customer ID (for widget users)
            user_id: Optional user ID (for authenticated users)
            allowed_types: Optional set of allowed MIME types for this upload
        
        Returns:
            Dict with 'file_url', 'signed_url', 'filename', 'content_type', 'size'
        
        Raises:
            ValueError: If validation fails or upload fails
        """
        try:
            filename = file_data.get('filename', '')
            content_type = file_data.get('content_type', 'application/octet-stream')
            
            # Validate file data
            if not file_data.get('content') or not filename:
                raise ValueError("Missing required file data: content and filename")
            
            # Determine which types are allowed
            types_allowed = allowed_types if allowed_types is not None else ALLOWED_FILE_TYPES
            friendly_types = FileUploadService.get_friendly_allowed_types_message(types_allowed)
            
            # Step 1: Validate extension matches claimed content type
            is_valid, error = FileUploadService.validate_extension(filename, content_type)
            if not is_valid:
                logger.warning(f"Extension validation failed for {filename}: {error}")
                raise ValueError(error)
            
            # Step 2: Validate file type is allowed
            if not FileUploadService.validate_file_type(content_type, types_allowed):
                raise ValueError(f"File type '{content_type}' is not allowed. Allowed types: {friendly_types}")
            
            # Step 3: Decode base64 content
            try:
                file_content = base64.b64decode(file_data['content'])
            except Exception as e:
                logger.error(f"Error decoding base64 content: {str(e)}")
                raise ValueError("Invalid file content encoding")
            
            file_size = len(file_content)
            
            # Step 4: Validate file size
            is_valid, error_message = FileUploadService.validate_file_size(file_size, content_type)
            if not is_valid:
                raise ValueError(error_message)
            
            # Step 5: SECURITY - Validate magic bytes to prevent MIME spoofing
            is_valid, magic_error = FileUploadService.validate_magic_bytes(file_content, content_type)
            if not is_valid:
                logger.warning(f"Magic byte validation failed for {filename}: {magic_error}")
                raise ValueError(f"Security check failed: {magic_error}")
            
            logger.info(f"File validation passed for {filename}: type={content_type}, size={file_size}")
            
            # Generate unique filename with sanitized extension
            file_extension = os.path.splitext(filename)[1].lower()
            # Additional sanitization: only allow known extensions
            if file_extension not in EXTENSION_TO_MIME:
                raise ValueError(f"File extension '{file_extension}' is not allowed")
            
            unique_filename = f"{uuid.uuid4()}{file_extension}"
            
            # Determine folder based on auth type
            folder = f"chat_attachments/{org_id}"
            
            # Upload to S3 if enabled, otherwise save locally
            if settings.S3_FILE_STORAGE:
                file_url = await upload_file_to_s3(
                    file_content=file_content,
                    folder=folder,
                    filename=unique_filename,
                    content_type=content_type
                )
            else:
                # Save locally
                upload_dir = os.path.join("uploads", folder)
                os.makedirs(upload_dir, exist_ok=True)
                
                file_path = os.path.join(upload_dir, unique_filename)
                with open(file_path, "wb") as f:
                    f.write(file_content)
                
                file_url = f"/uploads/{folder}/{unique_filename}"

            signed_url = await resolve_public_url(file_url)

            logger.info(f"Successfully uploaded file: {filename} -> {file_url}")
            
            return {
                'file_url': file_url,
                'signed_url': signed_url,
                'filename': filename,
                'content_type': content_type,
                'size': file_size
            }
            
        except ValueError:
            # Re-raise validation errors
            raise
        except Exception as e:
            logger.error(f"Error in FileUploadService.upload_file: {str(e)}", exc_info=True)
            raise ValueError(f"Failed to upload file: {str(e)}")