"""
AWS S3 IO Module

Utilities for reading and writing files from/to AWS S3

License: BSD 3-Clause
"""

import os
from pathlib import Path
from typing import Dict, Optional, Tuple, Any, Union
from urllib.parse import urlparse

import boto3
from botocore.client import BaseClient
from botocore.exceptions import NoCredentialsError, ClientError
from cocina.printer import Printer
from soundhub_utils.utils import flac
from soundhub_utils.io import _shared

# Setup logging
printer = Printer()

def extract_flac_header(src: str) -> Dict:
    """
    Extract FLAC header metadata by downloading only the first few KB.
    
    Args:
        src: S3 URI (may or may not include s3:// prefix)
        
    Returns:
        Dict containing duration, sample_rate, channels, total_samples
    """
    bucket, key = _parse_s3_uri(src)
    s3_client = boto3.client('s3')
    
    try:
        # Download first 4KB to get FLAC metadata
        response = s3_client.get_object(
            Bucket=bucket,
            Key=key,
            Range='bytes=0-4095'
        )
        
        header_data = response['Body'].read()
        
        # Use flac utility to parse header
        return flac.extract_header(header_data)
        
    except ClientError as e:
        printer.error(e, f"Error accessing S3 object {src}")
        raise
    except Exception as e:
        printer.error(e, f"Error parsing FLAC header from {src}")
        raise


def read(
    src: str,
    dest: Optional[str] = None,
    start_byte: Optional[int] = None,
    end_byte: Optional[int] = None,
    create_parents: bool = True,
    overwrite: bool = True,
    client: Optional[BaseClient] = None
) -> Union[str, Any]:
    """
    Read a file from S3 bucket.
    
    Args:
        src: S3 URI (may or may not include "s3://" prefix)
        dest: Local file destination (if None, return response from AWS download)
        start_byte: Start byte position for range request
        end_byte: End byte position for range request
        create_parents: Create parent directories when writing files
        overwrite: Overwrite existing destination file
        client: S3 client (if None, create one with boto3.client('s3'))
        
    Returns:
        str (destination file path) if dest is provided, otherwise response object from S3
    """
    if client is None:
        client = boto3.client('s3')
    
    bucket, key = _parse_s3_uri(src)
    
    try:
        # Prepare get_object parameters
        get_object_params = {
            'Bucket': bucket,
            'Key': key
        }
        
        # Add range request if specified
        if start_byte is not None or end_byte is not None:
            if start_byte is not None and end_byte is not None:
                get_object_params['Range'] = f'bytes={start_byte}-{end_byte}'
            elif start_byte is not None:
                get_object_params['Range'] = f'bytes={start_byte}-'
            elif end_byte is not None:
                get_object_params['Range'] = f'bytes=0-{end_byte}'
        
        # Get the object
        response = client.get_object(**get_object_params)
        
        # If no destination specified, return the response
        if dest is None:
            return response
        
        # Setup destination path
        dest_path = Path(dest)
        if create_parents:
            dest_path.parent.mkdir(parents=True, exist_ok=True)
        
        if dest_path.exists() and not overwrite:
            raise FileExistsError(f"Destination file {dest} already exists and overwrite=False")
        
        # Write to destination
        with open(dest_path, 'wb') as f:
            f.write(response['Body'].read())
        
        printer.message(f"Downloaded {src} to {dest_path}")
        return str(dest_path)
        
    except ClientError as e:
        printer.error(e, f"Error reading from S3 {src}")
        raise
    except Exception as e:
        printer.error(e, f"Error processing {src}")
        raise


def read_partial_flac(
    src: str,
    dest: str,
    start_time: float = None,
    end_time: Optional[float] = None,
    duration: Optional[float] = None,
    create_parents: bool = True,
    overwrite: bool = True
) -> str:
    """
    Read a part of a FLAC file from S3 bucket.
    
    Args:
        src: The S3 URI (may or may not include "s3://" prefix)
        dest: The local file destination
        start_time: Start time in seconds (default: 0)
        end_time: End time in seconds (if None, use start_time + duration)
        duration: Duration in seconds (ignored if end_time is provided)
        create_parents: Create parent directories if needed
        overwrite: Overwrite existing destination file
        
    Returns:
        Path to the saved file (dest)
    """
    # Validate inputs
    if start_time is None:
        start_time = 0.0
    
    if end_time is None:
        if duration is None:
            raise ValueError("Either end_time or duration must be provided")
        end_time = start_time + duration
    
    if start_time < 0:
        raise ValueError("start_time must be >= 0")
    
    if end_time <= start_time:
        raise ValueError("end_time must be > start_time")
    
    # Setup destination path
    dest_path = Path(dest)
    if create_parents:
        dest_path.parent.mkdir(parents=True, exist_ok=True)
    
    if dest_path.exists() and not overwrite:
        raise FileExistsError(f"Destination file {dest} already exists and overwrite=False")
    
    printer.message(f"Reading FLAC partial: {src} -> {dest}, time range: {start_time:.1f}-{end_time:.1f}s")

    try:
        # Get file metadata
        header_info = extract_flac_header(src)
        file_size = _get_file_size(src)
        total_duration = header_info['duration']

        # Validate time range
        if start_time >= total_duration:
            raise ValueError(f"start_time {start_time} >= file duration {total_duration}")

        end_time = min(end_time, total_duration)
        requested_duration = end_time - start_time

        printer.message(f"File info: duration={total_duration:.1f}s, size={file_size/(1024*1024):.1f}MB")

        # For partial requests, ALWAYS use byte-range downloads (never fall back to full file)
        start_byte, end_byte = flac.estimate_byte_range(
            file_size, total_duration, start_time, end_time
        )

        printer.message(f"Using byte-range download strategy: {start_byte}-{end_byte}")

        # Use shared function for byte-range download and extraction
        return _shared.download_and_extract_partial_flac(
            src, dest, start_time, end_time, requested_duration,
            start_byte, end_byte, create_parents, overwrite,
            _download_range
        )

    except Exception as e:
        printer.error(e, f"Error processing {src}")
        raise


#
# INTERNAL
#
def _parse_s3_uri(s3_uri: str) -> Tuple[str, str]:
    """Parse S3 URI into bucket and key components."""
    if s3_uri.startswith('s3://'):
        parsed = urlparse(s3_uri)
        return parsed.netloc, parsed.path.lstrip('/')
    else:
        # Assume format: bucket/key/path
        parts = s3_uri.split('/', 1)
        if len(parts) != 2:
            raise ValueError(f"Invalid S3 URI format: {s3_uri}")
        return parts[0], parts[1]


def _get_file_size(s3_uri: str) -> int:
    """Get file size from S3 without downloading."""
    bucket, key = _parse_s3_uri(s3_uri)
    s3_client = boto3.client('s3')
    
    try:
        response = s3_client.head_object(Bucket=bucket, Key=key)
        return response['ContentLength']
    except ClientError as e:
        printer.error(e, f"Error getting file size for {s3_uri}")
        raise


def _download_range(s3_uri: str, start_byte: int, end_byte: int) -> bytes:
    """Download a byte range from S3."""
    bucket, key = _parse_s3_uri(s3_uri)
    s3_client = boto3.client('s3')

    try:
        response = s3_client.get_object(
            Bucket=bucket,
            Key=key,
            Range=f'bytes={start_byte}-{end_byte}'
        )
        return response['Body'].read()
    except ClientError as e:
        printer.error(e, f"Error downloading range {start_byte}-{end_byte} from {s3_uri}")
        raise
