import logging
import shutil
import tempfile
from pathlib import Path

import numpy as np

from songmatch.audio_io import create_temp_wav, load_audio, resample, save_audio
from songmatch.config import SongMatchConfig
from songmatch.eq_match import match_eq
from songmatch.loudness import crop_reference
from songmatch.mix_match import run_matchering
from songmatch.separation import STEM_NAMES, separate_stems
from songmatch.utils import make_output_path, make_stem_dir

logger = logging.getLogger(__name__)


def run(config: SongMatchConfig) -> list[Path]:
    """Run the SongMatch pipeline. Returns list of output file paths produced."""
    outputs: list[Path] = []

    mix_match_input = config.user_audio

    if config.stems_match:
        stems_matched_path = _run_stems_match(config)
        outputs.append(stems_matched_path)
        mix_match_input = stems_matched_path

    if config.mix_match:
        mix_matched_path = _run_mix_match(config, target_audio=mix_match_input)
        outputs.append(mix_matched_path)

    return outputs


def _run_stems_match(config: SongMatchConfig) -> Path:
    """Execute the per-stem EQ matching pipeline."""
    # Step 1: Load and validate inputs
    logger.info("Loading user audio: %s", config.user_audio)
    user_audio, user_sr = load_audio(config.user_audio)

    logger.info("Loading reference audio: %s", config.reference_audio)
    ref_audio, ref_sr = load_audio(config.reference_audio)

    # Step 2: Loudest segment detection on reference
    logger.info("Analyzing loudness of reference...")
    cropped_ref, cropped_sr = crop_reference(ref_audio, ref_sr, config)

    # Step 3: Write cropped reference to temp file and separate stems
    temp_dir = Path(tempfile.mkdtemp(prefix="songmatch_"))
    try:
        cropped_ref_path = temp_dir / "ref_cropped.wav"
        save_audio(cropped_ref_path, cropped_ref, cropped_sr, bits=32)

        ref_stems_dir = temp_dir / "ref_stems"
        logger.info("Separating reference stems...")
        ref_stems = separate_stems(cropped_ref_path, ref_stems_dir, config.model)

        # Step 4: Separate user stems
        user_stem_dir = make_stem_dir(config.user_audio)
        logger.info("Separating user stems...")
        user_stems = separate_stems(config.user_audio, user_stem_dir, config.model)

        # Step 5: Per-stem EQ matching
        for stem_name in STEM_NAMES:
            if stem_name not in user_stems or stem_name not in ref_stems:
                logger.warning("Stem '%s' missing. Skipping EQ for this stem.", stem_name)
                continue

            logger.info("Matching EQ for %s...", stem_name)
            user_stem_data, stem_sr = load_audio(user_stems[stem_name])
            ref_stem_data, ref_stem_sr = load_audio(ref_stems[stem_name])

            # Ensure same sample rate
            if ref_stem_sr != stem_sr:
                ref_stem_data = resample(ref_stem_data, ref_stem_sr, stem_sr)

            eq_matched = match_eq(
                user_stem_data,
                ref_stem_data,
                stem_sr,
                config.eq_smoothing,
                config.eq_max_gain,
            )

            # Overwrite user stem with EQ-matched version
            save_audio(user_stems[stem_name], eq_matched, stem_sr, bits=32)

        # Step 6: Recombine stems
        logger.info("Recombining stems...")
        stems_matched_path = _recombine_stems(user_stems, config)

    finally:
        # Step 7: Cleanup
        if config.keep_ref_stems:
            logger.info("Keeping reference stems in %s", temp_dir)
        else:
            shutil.rmtree(temp_dir, ignore_errors=True)
            logger.debug("Cleaned up temp directory: %s", temp_dir)

    # Rename user stems to clean names
    _rename_user_stems(user_stems, make_stem_dir(config.user_audio))

    logger.info("Stems-match complete: %s", stems_matched_path)
    return stems_matched_path


def _recombine_stems(stems: dict[str, Path], config: SongMatchConfig) -> Path:
    """Sum EQ-matched stems into a single recombined file."""
    stem_arrays = []
    stem_sr = None

    for stem_name in STEM_NAMES:
        if stem_name not in stems:
            continue
        data, sr = load_audio(stems[stem_name])
        if stem_sr is None:
            stem_sr = sr
        elif sr != stem_sr:
            data = resample(data, sr, stem_sr)
        stem_arrays.append(data)

    if not stem_arrays:
        raise RuntimeError("No stems to recombine.")

    # Pad shorter stems with zeros
    max_len = max(a.shape[0] for a in stem_arrays)
    max_channels = max(a.shape[1] if a.ndim == 2 else 1 for a in stem_arrays)

    padded = []
    for a in stem_arrays:
        if a.ndim == 1:
            a = a[:, np.newaxis]
        if a.shape[1] < max_channels:
            a = np.tile(a, (1, max_channels))[:, :max_channels]
        if a.shape[0] < max_len:
            padding = np.zeros((max_len - a.shape[0], a.shape[1]), dtype=np.float32)
            a = np.concatenate([a, padding], axis=0)
        padded.append(a)

    # Sum all stems
    combined = np.sum(padded, axis=0).astype(np.float32)

    # Peak normalize if clipping
    peak = np.max(np.abs(combined))
    if peak > 1.0:
        target = 10 ** (-0.1 / 20)  # -0.1 dBFS
        combined = combined * (target / peak)
        logger.info("Peak normalized recombined audio (peak was %.2f)", peak)

    # Resample to output sample rate if needed
    if stem_sr != config.output_sr:
        combined = resample(combined, stem_sr, config.output_sr)

    output_path = make_output_path(config.user_audio, "stemsmatched")
    save_audio(output_path, combined, config.output_sr, config.output_bits)
    return output_path


def _rename_user_stems(stems: dict[str, Path], stem_dir: Path) -> None:
    """Rename separated stem files to clean names (vocals.wav, drums.wav, etc.)."""
    for stem_name, stem_path in stems.items():
        target = stem_dir / f"{stem_name}.wav"
        if stem_path != target and stem_path.exists():
            shutil.move(str(stem_path), str(target))


def _run_mix_match(config: SongMatchConfig, target_audio: Path) -> Path:
    """Execute the matchering pipeline."""
    output_path = make_output_path(config.user_audio, "mixmatched")

    logger.info("Running mix-match...")
    run_matchering(
        target_audio=target_audio,
        reference_audio=config.reference_audio,
        output_path=output_path,
        config=config,
    )

    logger.info("Mix-match complete: %s", output_path)
    return output_path
