"""
TensorFlow Utilities Module

Generic TensorFlow utilities for working with TFRecords and other TF-specific operations.

License: BSD 3-Clause
"""
import os
from pathlib import Path
from typing import Union, Optional, List, Dict, Any
import numpy as np
import tensorflow as tf


class TFRecordWriter:
    """
    Generic context manager for streaming TFRecord writes.

    Writes data incrementally to avoid memory issues with large datasets.
    Supports optional chunking to split into multiple TFRecord files.

    Example:
        with TFRecordWriter('output.tfrecord', compression='GZIP') as writer:
            for i, data in enumerate(dataset):
                features = {'data': tf.train.Feature(...)}
                writer.write_example(features)
    """

    def __init__(
        self,
        output_path: Union[str, Path],
        compression: Optional[str] = 'GZIP',
        chunk_size: Optional[int] = None,
        basename: str = 'records'
    ):
        """
        Initialize TFRecord writer.

        Args:
            output_path: Path to output TFRecord file (or directory if chunking)
            compression: Compression type ('GZIP', 'ZLIB', or None)
            chunk_size: Number of records per TFRecord file (None = single file)
            basename: Base name for chunked files (e.g., 'data_001.tfrecord')
        """
        self.output_path = Path(output_path)
        self.compression = compression
        self.chunk_size = chunk_size
        self.basename = basename
        self.writer = None
        self.current_chunk = 0
        self.chunk_count = 0
        self.written_files = []

    def __enter__(self):
        # Create parent directories
        if self.chunk_size is None:
            # Single file mode
            self.output_path.parent.mkdir(parents=True, exist_ok=True)
            self._open_writer(self.output_path)
        else:
            # Chunked mode - output_path is directory
            self.output_path.mkdir(parents=True, exist_ok=True)

        return self

    def _open_writer(self, path: Path):
        """Open a new TFRecord writer."""
        options = None
        if self.compression:
            options = tf.io.TFRecordOptions(compression_type=self.compression)

        self.writer = tf.io.TFRecordWriter(str(path), options=options)
        self.written_files.append(str(path))

    def _close_writer(self):
        """Close current TFRecord writer."""
        if self.writer:
            self.writer.close()
            self.writer = None

    def _maybe_rotate_chunk(self):
        """Check if we need to start a new chunk file."""
        if self.chunk_size is None:
            return  # Single file mode

        if self.chunk_count >= self.chunk_size:
            # Close current writer and open new one
            self._close_writer()
            self.current_chunk += 1
            self.chunk_count = 0

            # Generate new filename
            chunk_filename = f"{self.basename}_chunk{self.current_chunk:04d}.tfrecord"
            chunk_path = self.output_path / chunk_filename
            self._open_writer(chunk_path)

    def write_example(self, features: Dict[str, 'tf.train.Feature']):
        """
        Write a single example to TFRecord.

        Args:
            features: Dictionary of TensorFlow features
        """
        # Check if we need to open first writer (chunked mode)
        if self.chunk_size is not None and self.writer is None:
            chunk_filename = f"{self.basename}_chunk{self.current_chunk:04d}.tfrecord"
            chunk_path = self.output_path / chunk_filename
            self._open_writer(chunk_path)

        # Check if we need to rotate to new chunk
        self._maybe_rotate_chunk()

        # Create and write example
        example = tf.train.Example(features=tf.train.Features(feature=features))
        self.writer.write(example.SerializeToString())
        self.chunk_count += 1

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._close_writer()

    def get_written_files(self) -> List[str]:
        """Get list of all TFRecord files written."""
        return self.written_files


class TFRecordSpectrogramWriter:
    """
    Specialized TFRecord writer for audio spectrograms.

    Uses the generic TFRecordWriter and adds spectrogram-specific functionality.

    Example:
        with TFRecordSpectrogramWriter('output.tfrecord', compression='GZIP') as writer:
            for i, spec in enumerate(spectrograms):
                writer.write_spectrogram(spec, i, 'audio_file')
    """

    def __init__(
        self,
        output_path: Union[str, Path],
        compression: Optional[str] = 'GZIP',
        chunk_size: Optional[int] = None,
        basename: str = 'spectrograms'
    ):
        """
        Initialize spectrogram TFRecord writer.

        Args:
            output_path: Path to output TFRecord file (or directory if chunking)
            compression: Compression type ('GZIP', 'ZLIB', or None)
            chunk_size: Number of spectrograms per TFRecord file (None = single file)
            basename: Base name for chunked files (e.g., 'spectrograms_001.tfrecord')
        """
        self.writer = TFRecordWriter(output_path, compression, chunk_size, basename)

    def __enter__(self):
        self.writer.__enter__()
        return self

    def write_spectrogram(
        self,
        spectrogram: np.ndarray,
        segment_index: int,
        filename: str,
        metadata: Optional[Dict[str, Any]] = None
    ):
        """
        Write a single spectrogram to TFRecord.

        Args:
            spectrogram: Spectrogram array (uint8, shape: [height, width])
            segment_index: Segment index in the audio file
            filename: Original audio filename (for tracking)
            metadata: Optional metadata dict (e.g., {'audio_uri': 's3://...', 'start_time': 500})
        """
        height, width = spectrogram.shape

        # Build base feature dictionary
        features = {
            'spectrogram': tf.train.Feature(
                bytes_list=tf.train.BytesList(value=[spectrogram.tobytes()])
            ),
            'height': tf.train.Feature(
                int64_list=tf.train.Int64List(value=[height])
            ),
            'width': tf.train.Feature(
                int64_list=tf.train.Int64List(value=[width])
            ),
            'segment_index': tf.train.Feature(
                int64_list=tf.train.Int64List(value=[segment_index])
            ),
            'filename': tf.train.Feature(
                bytes_list=tf.train.BytesList(value=[filename.encode('utf-8')])
            ),
        }

        # Add metadata fields if provided
        if metadata:
            for key, value in metadata.items():
                if isinstance(value, (int, bool)):
                    features[key] = tf.train.Feature(
                        int64_list=tf.train.Int64List(value=[int(value)])
                    )
                elif isinstance(value, float):
                    features[key] = tf.train.Feature(
                        float_list=tf.train.FloatList(value=[value])
                    )
                elif isinstance(value, str):
                    features[key] = tf.train.Feature(
                        bytes_list=tf.train.BytesList(value=[value.encode('utf-8')])
                    )
                elif value is None:
                    # Skip None values
                    continue
                else:
                    # Convert to string for unsupported types
                    features[key] = tf.train.Feature(
                        bytes_list=tf.train.BytesList(value=[str(value).encode('utf-8')])
                    )

        # Write using generic writer
        self.writer.write_example(features)

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.writer.__exit__(exc_type, exc_val, exc_tb)

    def get_written_files(self) -> List[str]:
        """Get list of all TFRecord files written."""
        return self.writer.get_written_files()


def load_spectrograms_from_tfrecord(
    tfrecord_paths: Union[str, List[str]],
    spectrogram_shape: tuple[int, int],
    batch_size: int = 100,
    compression: Optional[str] = 'GZIP'
):
    """
    Load spectrograms from TFRecord file(s) for inference.

    Args:
        tfrecord_paths: Path to TFRecord file or list of TFRecord files (for chunked mode)
        spectrogram_shape: Expected spectrogram shape (height, width)
        batch_size: Batch size for dataset
        compression: Compression type used ('GZIP', 'ZLIB', or None)

    Returns:
        TensorFlow dataset yielding batches of (spectrograms, metadata_dict)
        - spectrograms: shape (batch_size, height, width, 1), float32, range [0, 1]
        - metadata_dict: dict with 'segment_index', 'filename', and any custom metadata

    Example:
        dataset = load_spectrograms_from_tfrecord('audio.tfrecord', (257, 1000))
        for spectrograms, metadata in dataset:
            predictions = model.predict(spectrograms)
    """
    # Ensure tfrecord_paths is a list
    if isinstance(tfrecord_paths, str):
        tfrecord_paths = [tfrecord_paths]

    def parse_example(example_proto):
        """Parse a single TFRecord example."""
        # Define base feature description
        feature_description = {
            'spectrogram': tf.io.FixedLenFeature([], tf.string),
            'height': tf.io.FixedLenFeature([], tf.int64),
            'width': tf.io.FixedLenFeature([], tf.int64),
            'segment_index': tf.io.FixedLenFeature([], tf.int64),
            'filename': tf.io.FixedLenFeature([], tf.string),
        }

        parsed = tf.io.parse_single_example(example_proto, feature_description)

        # Decode spectrogram
        spectrogram = tf.io.decode_raw(parsed['spectrogram'], tf.uint8)
        spectrogram = tf.reshape(spectrogram, [spectrogram_shape[0], spectrogram_shape[1]])
        spectrogram = tf.cast(spectrogram, tf.float32) / 255.0
        spectrogram = tf.expand_dims(spectrogram, -1)  # Add channel dimension

        # Build metadata dictionary
        metadata = {
            'segment_index': parsed['segment_index'],
            'filename': parsed['filename'],
            'height': parsed['height'],
            'width': parsed['width'],
        }

        return spectrogram, metadata

    # Create dataset from TFRecord file(s)
    dataset = tf.data.TFRecordDataset(
        tfrecord_paths,
        compression_type=compression if compression else ''
    )
    dataset = dataset.map(parse_example, num_parallel_calls=tf.data.AUTOTUNE)
    dataset = dataset.batch(batch_size)
    dataset = dataset.prefetch(tf.data.AUTOTUNE)  # Prefetch for better performance

    return dataset


def load_spectrograms_tf(file_paths: List[str], spectrogram_shape: tuple[int, int]) -> 'tf.Dataset':
    """
    Load spectrogram images using TensorFlow for model compatibility.

    Args:
        file_paths: List of paths to spectrogram PNG files
        spectrogram_shape: Target shape (height, width)

    Returns:
        TensorFlow dataset with batched images
    """
    def load_and_preprocess_image(path):
        image = tf.io.read_file(path)
        image = tf.image.decode_png(image, channels=1)
        image = tf.cast(image, tf.float32) / 255.0

        # Don't resize - images should already be correct size
        # Resizing causes interpolation artifacts that affect predictions
        return image

    # Create dataset from file paths and load images
    path_ds = tf.data.Dataset.from_tensor_slices(file_paths)
    image_ds = path_ds.map(load_and_preprocess_image, num_parallel_calls=tf.data.AUTOTUNE)

    # Batch all images
    batch_ds = image_ds.batch(100)

    return batch_ds


def is_tfrecord_path(path: Union[str, Path, List[str], List[Path]]) -> bool:
    """
    Check if path(s) point to TFRecord file(s).

    Args:
        path: File path, list of paths, or directory path

    Returns:
        True if path(s) are TFRecord files, False otherwise
    """
    # Handle list of paths
    if isinstance(path, list):
        if len(path) == 0:
            return False
        # Check first item
        path = path[0]

    path = Path(path)

    # Check extension
    return path.suffix.lower() == '.tfrecord' or path.name.endswith('.tfrecord')