"""
Audio to spectrogram Script

FIX ME
"""
from cocina.printer import Printer
from soundhub_utils.io import io
printer = Printer()


def get_audio():
    if start_time is not None or end_time is not None:
        # Partial download
        download_duration = None
        if start_time is not None and end_time is not None:
            download_duration = end_time - start_time
        elif end_time is not None:
            start_time = 0.0
            download_duration = end_time
        
        io.read_partial_flac(
            src=source_path,
            dest=temp_audio_path,
            start_time=start_time if start_time is not None else 0.0,
            duration=download_duration,
            overwrite=True
        )
        
        actual_start_time = start_time if start_time is not None else 0.0
        printer.message(f"Downloaded partial audio: {actual_start_time}s to {actual_start_time + (download_duration or 0)}s")
        
        # Validate the downloaded file has content
        file_size = os.path.getsize(temp_audio_path) if os.path.exists(temp_audio_path) else 0
        if file_size < 1000:  # Less than 1KB
            raise IOError(f"Partial FLAC extraction failed - output file too small: {file_size} bytes. "
                        f"This is likely due to an issue with the byte-range download and time extraction process. "
                        f"Try using a full file download or different time range.")
    else:
        # Full file download
        io.read(
            src=source_path,
            dest=temp_audio_path,
            overwrite=True
        )
        actual_start_time = 0.0
        printer.message("Downloaded full audio file")


def preprocess():
	pass


def spectorgrams():
    if keep_spectrograms:
        if spectrogram_dir is None:
            spectrogram_dir = output_path.parent / f"{output_path.stem}_spectrograms"
        else:
            spectrogram_dir = Path(spectrogram_dir)
        
        spectrogram_dir.mkdir(parents=True, exist_ok=True)
        printer.message(f"Spectrograms will be saved to: {spectrogram_dir}")
        
        # Preprocess and save spectrograms
        spectrogram_paths = model_owl_dev.preprocess(
            audio_file=temp_audio_path,
            n_segments=n_segments,
            overlap=overlap,
            save_spectrograms=True,
            output_dir=str(spectrogram_dir),
            overwrite=True,
            segment_duration=segment_duration,
            target_sample_rate=TARGET_SAMPLE_RATE,
            spectrogram_shape=SPECTROGRAM_SHAPE
        )
        
        spectrograms_input = spectrogram_paths
        printer.message(f"Saved {len(spectrogram_paths)} spectrogram files")
    else:
        # Process spectrograms in memory
        spectrograms = model_owl_dev.preprocess(
            audio_file=temp_audio_path,
            n_segments=n_segments,
            overlap=overlap,
            save_spectrograms=False,
            segment_duration=segment_duration,
            target_sample_rate=TARGET_SAMPLE_RATE,
            spectrogram_shape=SPECTROGRAM_SHAPE
        )
        
        spectrograms_input = spectrograms
        printer.message(f"Generated {len(spectrograms)} spectrogram arrays in memory")


def inference():
    # Step 4: Calculate segment start times for extended parquet format
    segment_start_times = []
    for i in range(n_segments):
        segment_start_times.append(actual_start_time + (i * (segment_duration - overlap)))
    
    # Step 5: Run inference with extended parquet format
    results_path = model_owl_dev.run_inference(
        spectrograms=spectrograms_input,
        model_path=str(model_path),
        return_format='parquet',
        output_file=str(output_path),
        overwrite=True,
        spectrogram_shape=SPECTROGRAM_SHAPE,
        class_names=CLASS_NAMES,
        flac_filename=flac_filename,
        segment_start_times=segment_start_times,
        segment_duration=segment_duration,
        processing_timestamp=datetime.now().isoformat()
    )
    
    end_timestamp = datetime.now()
    duration = end_timestamp - start_timestamp



















def preprocess(
    audio_file: str,
    n_segments: int = 1,
    overlap: float = 0.0,
    save_spectrograms: bool = False,
    output_dir: Optional[str] = None,
    create_parents: bool = True,
    overwrite: bool = True,
    segment_duration: float = _DEFAULT_SEGMENT_DURATION,
    target_sample_rate: int = _DEFAULT_TARGET_SAMPLE_RATE,
    spectrogram_shape: Tuple[int, int] = _DEFAULT_SPECTROGRAM_SHAPE
) -> Union[List[np.ndarray], List[str]]:
    """
    Preprocess audio file into model-ready spectrograms.
    
    Args:
        audio_file: Path to audio file (already downloaded)
        n_segments: Number of segments to extract
        overlap: Overlap between segments in seconds
        save_spectrograms: If True, save spectrograms to files and return paths
        output_dir: Directory to save spectrograms (if save_spectrograms=True)
        create_parents: Create parent directories if needed
        overwrite: Overwrite existing files
        segment_duration: Duration of each segment in seconds
        target_sample_rate: Target sample rate for audio processing
        spectrogram_shape: Expected spectrogram shape (height, width)
        
    Returns:
        List of spectrogram numpy arrays if save_spectrograms=False,
        otherwise list of file paths to saved PNG spectrograms
        
    Raises:
        ValueError: If invalid parameters provided
        IOError: If audio processing fails
    """
    if overlap >= segment_duration:
        raise ValueError("Overlap must be less than segment duration")
    
    printer.message(f"Preprocessing {n_segments} segments from {audio_file}")
    
    try:
        # Load and resample audio using soundhub_utils.utils.audio
        audio_array, original_sr = audio.pydub_audio_to_numpy(audio_file)
        printer.message(f"Loaded audio: {len(audio_array)} samples at {original_sr}Hz")
        
        # Resample to target sample rate if needed
        if original_sr != target_sample_rate:
            num_samples = int(len(audio_array) * target_sample_rate / original_sr)
            audio_array = resample(audio_array, num_samples)
            sample_rate = target_sample_rate
            printer.message(f"Resampled from {original_sr}Hz to {sample_rate}Hz")
        else:
            sample_rate = original_sr
        
        # Extract segments and create spectrograms
        spectrograms = []
        spectrogram_paths = []
        
        if save_spectrograms and output_dir:
            output_path = Path(output_dir)
            if create_parents:
                output_path.mkdir(parents=True, exist_ok=True)
        
        for i in range(n_segments):
            # Calculate segment boundaries
            segment_start_time = i * (segment_duration - overlap)
            segment_start_sample = int(segment_start_time * sample_rate)
            segment_end_sample = segment_start_sample + int(segment_duration * sample_rate)
            
            if segment_end_sample > len(audio_array):
                logger.warning(f"Segment {i+1} extends beyond audio, skipping")
                break
            
            segment_audio = audio_array[segment_start_sample:segment_end_sample]
            printer.message(f"Processing segment {i+1}: {segment_start_time:.1f}s - {segment_start_time + segment_duration:.1f}s")
            
            # Create spectrogram
            spectrogram = _create_spectrogram(segment_audio, sample_rate, spectrogram_shape)
            
            if save_spectrograms and output_dir:
                # Save as PNG file
                filename = f'spectrogram_segment_{i+1:02d}.png'
                filepath = output_path / filename
                _save_spectrogram_png(spectrogram, str(filepath), overwrite, spectrogram_shape)
                spectrogram_paths.append(str(filepath))
            else:
                # Keep as numpy array
                spectrograms.append(spectrogram)
        
        return spectrogram_paths if save_spectrograms else spectrograms
        
    except Exception as e:
        printer.error(e, f"Preprocessing failed")
        raise IOError(f"Failed to preprocess audio from {audio_file}: {e}")
