import os
import re
import json
import shutil
import random
import subprocess
import pandas as pd
from tqdm import tqdm
from pathlib import Path
from typing import List, Dict
from joblib import Parallel, delayed
from suno_utils.utils.text import read_jsonl

# -------------------------
# Configuration
# -------------------------

# t<n> - data collection
# ab<n> - ab test

EXPERIMENT_NAME = "t5"
# PARENT_DIR = "/Users/christiansteinmetz/Downloads/v3-distill-data-ctx-t1"  # <-- root with UUID dirs
PARENT_DIR = "/app2/suno/data/christian/outputs/v3-base-data-ctx-t2"
# PARENT_DIR = "/app2/suno/data/christian/outputs/v3-base-data-ctx-t1"

OUTPUT_DIR = f"/app2/suno/data/christian/outputs/labelmaker-ab/{EXPERIMENT_NAME}"
# METADATA_DIR = f"/home/christian/code/christian/metadata/labelmaker/{EXPERIMENT_NAME}"
METADATA_DIR = f"{PARENT_DIR}/metadata"

# Mode selection: "perfect_pairs" or "model_comparison"
COMPARISON_MODE = "perfect_pairs"  # Change to "perfect_pairs" for original behavior

# Example configurations:
#
# Perfect pairs mode (original behavior):
# COMPARISON_MODE = "perfect_pairs"
# MODEL_SUBSTRINGS = []  # Empty for all models, or specify to filter
#
# Model comparison mode (new behavior):
# COMPARISON_MODE = "model_comparison"
# MODEL_SUBSTRINGS = ["model_a", "model_b"]  # Exactly 2 models to compare

MODEL_SUBSTRINGS: List[str] = [
    "4n_25hz_2b_flow_5e5_sft_t8_500k",
    # "v3_flow_distill_s1039_v1_t18_1E6_beta100_n15_bt2_acc2_4k_last_step2",
    # "v2_infill_d4_t39_1E6_beta100_n16_bt2_acc4_3k_last_step10",
]  # For model_comparison mode, these are the two models to compare
# For perfect_pairs mode, leave empty for all models or specify to filter

# Audio processing settings for model comparison mode
CROP_DURATION = 30.0  # seconds to crop audio to
TARGET_LOUDNESS = -16.0  # target loudness in dB for normalization

UPLOAD_S3 = True
N_JOBS_S3 = 32  # number of parallel upload workers

INTERNAL_S3_BUCKET = f"s3://suno-annotation-public/christian/{EXPERIMENT_NAME}"
PUBLIC_S3_HTTP_BASE = (
    f"https://suno-annotation-public.s3.amazonaws.com/christian/{EXPERIMENT_NAME}"
)

PAIR_RE = re.compile(r"^(?P<base>.+?)[_-](?P<idx>[0-9])\.mp3$", re.IGNORECASE)

HONEYPOT_COUNT = 100  # number of honeypot pairs
HONEYPOT_CUTOFF_MIN_HZ = 2000  # Hz cutoff for lowpass filter
HONEYPOT_CUTOFF_MAX_HZ = 4000  # Hz cutoff for lowpass filter
RANDOM_SEED = 0  # reproducible honeypot selection

# -------------------------
# Helpers
# -------------------------


def match_model(filename: str, substrings: List[str]) -> bool:
    if not substrings:
        return True
    return any(s in filename for s in substrings)


def parse_pair(filename: str):
    m = PAIR_RE.match(filename)
    if not m:
        return None, None
    return m.group("base"), m.group("idx")


def apply_lowpass_inplace(path: Path, cutoff: int):
    tmp_path = path.with_suffix(".tmp.mp3")
    cmd = [
        "ffmpeg",
        "-y",
        "-i",
        str(path),
        "-af",
        f"lowpass=f={cutoff}",
        "-c:a",
        "libmp3lame",
        "-b:a",
        "192k",
        str(tmp_path),
    ]
    subprocess.run(
        cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True
    )
    tmp_path.replace(path)


def aws_s3_cp(local: Path, s3_uri: str):
    subprocess.run(["aws", "s3", "cp", str(local), s3_uri], check=False)


def crop_and_normalize_audio_ffmpeg(
    input_path: Path,
    output_path: Path,
    duration: float = 30.0,
    target_loudness: float = -16.0,
):
    """Crop audio to random segment and normalize loudness using ffmpeg."""
    try:
        # Get audio duration
        duration_cmd = [
            "ffprobe",
            "-v",
            "quiet",
            "-show_entries",
            "format=duration",
            "-of",
            "csv=p=0",
            str(input_path),
        ]
        result = subprocess.run(
            duration_cmd, capture_output=True, text=True, check=True
        )
        total_duration = float(result.stdout.strip())

        # Calculate random start time
        if total_duration <= duration:
            # If audio is shorter than target duration, use the whole thing
            start_time = 0
            crop_duration = total_duration
        else:
            # Randomly select a segment
            max_start = total_duration - duration
            start_time = random.uniform(0, max_start)
            crop_duration = duration

        # Build ffmpeg command for cropping and loudness normalization
        cmd = [
            "ffmpeg",
            "-y",
            "-i",
            str(input_path),
            "-ss",
            str(start_time),
            "-t",
            str(crop_duration),
            "-af",
            f"loudnorm=I={target_loudness}:LRA=11:TP=-1.5",
            "-c:a",
            "libmp3lame",
            "-b:a",
            "192k",
            str(output_path),
        ]

        subprocess.run(
            cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True
        )

    except Exception as e:
        print(f"Error processing {input_path}: {e}")
        # Fallback to copying original file
        shutil.copy2(input_path, output_path)


def extract_model_name(filename: str) -> str:
    """Extract model name from filename based on comparison mode."""
    if COMPARISON_MODE == "perfect_pairs":
        # For perfect pairs, try to extract model name from the filename
        # Remove the index suffix and .mp3 extension using regex
        base, idx = parse_pair(filename)
        if base is not None:
            base_name = base
        else:
            # Fallback: remove .mp3 extension
            base_name = filename.replace(".mp3", "")

        # Try to extract a meaningful model identifier
        if "_" in base_name:
            # Take the last part as model identifier
            return base_name.split("_")[-1]
        else:
            return "unknown_model"

    elif COMPARISON_MODE == "model_comparison":
        # For model comparison, find which model substring matches
        for model_substring in MODEL_SUBSTRINGS:
            if model_substring in filename:
                # Return the full model substring
                return model_substring
        return "unknown_model"

    return "unknown_model"


def validate_config():
    """Validate configuration based on comparison mode."""
    if COMPARISON_MODE not in ["perfect_pairs", "model_comparison"]:
        raise ValueError(f"Unknown COMPARISON_MODE: {COMPARISON_MODE}")

    if COMPARISON_MODE == "model_comparison" and len(MODEL_SUBSTRINGS) != 2:
        raise ValueError("MODEL_COMPARISON mode requires exactly 2 model substrings")

    print(f"Configuration validated:")
    print(f"  Mode: {COMPARISON_MODE}")
    print(f"  Models: {MODEL_SUBSTRINGS}")
    print(f"  Parent dir: {PARENT_DIR}")
    print(f"  Output dir: {OUTPUT_DIR}")


# -------------------------
# Main
# -------------------------


def main():
    random.seed(RANDOM_SEED)
    validate_config()

    parent = Path(PARENT_DIR)
    out_base = Path(OUTPUT_DIR)
    metadata_dir = Path(METADATA_DIR)
    out_base.mkdir(parents=True, exist_ok=True)

    manifest = []
    all_pairs = []

    # check if there is a filtered_meta_pairs.jsonl file in the parent directory
    filtered_meta_pairs_path = metadata_dir / "filtered_meta_pairs.jsonl"
    if filtered_meta_pairs_path.exists():
        print(f"Found filtered meta pairs file: {filtered_meta_pairs_path}")
        filtered_meta_pairs = read_jsonl(filtered_meta_pairs_path)
        # convert to a dict with dirname as the key
        filtered_meta_pairs = {
            meta_pair["dirname"]: meta_pair for meta_pair in filtered_meta_pairs
        }
        print(f"Filtered meta pairs: {len(filtered_meta_pairs)}")
        # create a set of ids that are in the filtered_meta_pairs
        filtered_meta_pairs_ids = set(filtered_meta_pairs.keys())
    else:
        print(f"No filtered meta pairs file found: {filtered_meta_pairs_path}")
        filtered_meta_pairs = None
        filtered_meta_pairs_ids = set()

    if COMPARISON_MODE == "perfect_pairs":
        # Step 1: Gather all complete pairs (original behavior)
        for uuid_dir in tqdm(sorted(parent.iterdir())):
            if not uuid_dir.is_dir():
                continue

            if filtered_meta_pairs is not None:
                if uuid_dir.name not in filtered_meta_pairs_ids:
                    continue
                index_0 = filtered_meta_pairs[uuid_dir.name]["idx_pair"][0]
                index_1 = filtered_meta_pairs[uuid_dir.name]["idx_pair"][1]
            else:
                index_0 = 0
                index_1 = 1

            # Instead of iterating, construct the expected mp3 filenames for the two indices and check if they exist
            # We'll assume the naming convention is consistent with what parse_pair expects.
            # To do this, we need to find a sample file to extract the base name format.
            mp3_file_0 = (
                uuid_dir / f"{uuid_dir.name}_{MODEL_SUBSTRINGS[0]}_{index_0}.mp3"
            )
            mp3_file_1 = (
                uuid_dir / f"{uuid_dir.name}_{MODEL_SUBSTRINGS[0]}_{index_1}.mp3"
            )
            if mp3_file_0.exists() and mp3_file_1.exists():
                all_pairs.append((uuid_dir.name, mp3_file_0, mp3_file_1))

    elif COMPARISON_MODE == "model_comparison":
        # Step 1: Gather pairs from two different models
        if len(MODEL_SUBSTRINGS) != 2:
            raise ValueError(
                "MODEL_COMPARISON mode requires exactly 2 model substrings"
            )

        model0, model1 = MODEL_SUBSTRINGS[0], MODEL_SUBSTRINGS[1]

        for uuid_dir in sorted(parent.iterdir()):
            if not uuid_dir.is_dir():
                continue

            # Find files for each model
            model0_files = []
            model1_files = []

            for f in uuid_dir.iterdir():
                if not f.is_file() or f.suffix.lower() != ".mp3":
                    continue

                if model0 in f.name:
                    model0_files.append(f)
                elif model1 in f.name:
                    model1_files.append(f)

            # Create pairs from matching base names
            for f0 in model0_files:
                base0, _ = parse_pair(f0.name)
                if base0 is None:
                    continue

                # Find matching file from model1 with same base
                for f1 in model1_files:
                    base1, _ = parse_pair(f1.name)
                    if base1 == base0:
                        all_pairs.append((uuid_dir.name, f0, f1))
                        break

            # If no pairs found with exact base match, try alternative matching
            if not any(pair[0] == uuid_dir.name for pair in all_pairs):
                # Try matching by removing model names and comparing
                for f0 in model0_files:
                    # Remove model name from filename to get base
                    base0 = (
                        f0.name.replace(model0, "").replace("_", "").replace(".mp3", "")
                    )

                    for f1 in model1_files:
                        base1 = (
                            f1.name.replace(model1, "")
                            .replace("_", "")
                            .replace(".mp3", "")
                        )

                        if base0 == base1:
                            all_pairs.append((uuid_dir.name, f0, f1))
                            break

    else:
        raise ValueError(f"Unknown COMPARISON_MODE: {COMPARISON_MODE}")

    if not all_pairs:
        print("No complete pairs found.")
        return

    print(f"Found {len(all_pairs)} pairs in {COMPARISON_MODE} mode")
    if all_pairs:
        print(
            f"Sample pair: {all_pairs[0][0]} -> {all_pairs[0][1].name} vs {all_pairs[0][2].name}"
        )

    # Step 2: Pick honeypot pairs
    honeypot_pairs = set(random.sample(all_pairs, min(HONEYPOT_COUNT, len(all_pairs))))
    print(f"Honeypot pairs: {honeypot_pairs}")

    # Step 3: Process audio + build manifest (no S3 yet)
    upload_jobs = []

    for uuid, file0, file1 in tqdm(all_pairs):
        uuid_out_dir = out_base / uuid
        uuid_out_dir.mkdir(parents=True, exist_ok=True)

        copied_items = []
        pair_is_honeypot = (uuid, file0, file1) in honeypot_pairs
        bad_idx = random.choice([0, 1]) if pair_is_honeypot else None

        # Randomly swap the order of file0 and file1 for this pair
        files = [file0, file1]
        # order = [0, 1]
        # random.shuffle(order)
        # files = [files[i] for i in order]

        # If this is a honeypot, adjust bad_idx to match the new order
        # if pair_is_honeypot:
        #    # bad_idx refers to the original [file0, file1] order
        #    # Find where that file ended up after shuffling
        #    bad_idx_shuffled = order.index(bad_idx)
        # else:
        #    bad_idx_shuffled = None

        for idx, src in enumerate(files):
            dst = uuid_out_dir / src.name

            # Apply audio processing for model comparison mode
            if COMPARISON_MODE == "model_comparison":
                # For model comparison, crop and normalize audio
                crop_and_normalize_audio_ffmpeg(
                    src, dst, CROP_DURATION, TARGET_LOUDNESS
                )
            else:
                # For perfect pairs, just copy the file
                shutil.copy2(src, dst)

            if pair_is_honeypot and idx == bad_idx:
                apply_lowpass_inplace(
                    dst, random.randint(HONEYPOT_CUTOFF_MIN_HZ, HONEYPOT_CUTOFF_MAX_HZ)
                )
                source_tag = "honeypot_bad"
            elif pair_is_honeypot:
                source_tag = "honeypot_good"
            else:
                source_tag = None

            # Public name without model info
            file_index = src.name.split("_")[-1].replace(".mp3", "")
            sanitized_key = f"{uuid}-{file_index}.mp3"
            s3_uri = f"{INTERNAL_S3_BUCKET}/{sanitized_key}"
            public_url = f"{PUBLIC_S3_HTTP_BASE}/{sanitized_key}"

            upload_jobs.append((dst, s3_uri))

            item_meta = {
                "uuid": uuid,
                "model": extract_model_name(src.name),
                "file_index": file_index,
            }
            if source_tag:
                item_meta["source"] = source_tag

            letter = "A" if idx == 0 else "B"

            copied_items.append(
                {
                    "key": f"{uuid}_{idx}",
                    "name": f"Clip {letter}",
                    "url": public_url,
                    "metadata": item_meta,
                }
            )

        manifest.append(copied_items)

    # Step 4: Save manifest locally
    manifest_path = Path(OUTPUT_DIR) / "pairs.json"
    with open(manifest_path, "w") as f:
        json.dump(manifest, f, indent=2)
    print(f"Wrote manifest with {len(manifest)} pairs to: {manifest_path}")
    print(f"Honeypot pairs created: {len(honeypot_pairs)}")
    print(f"Mode: {COMPARISON_MODE}")

    # Step 5: Parallel S3 upload
    if UPLOAD_S3:
        print(f"Uploading {len(upload_jobs)} files to S3 with {N_JOBS_S3} workers...")
        Parallel(n_jobs=N_JOBS_S3)(
            delayed(aws_s3_cp)(local, s3_uri) for local, s3_uri in upload_jobs
        )
        print("S3 upload complete.")


if __name__ == "__main__":
    main()