import os
import json
from tqdm import tqdm


localid = int(os.environ.get("SLURM_LOCALID", 0))

os.environ["CUDA_VISIBLE_DEVICES"] = str(localid)

podcast_dir = "/app2/suno/data/podcast/"
podcast_meta_file = os.path.join(podcast_dir, "stratified_metadata_v2.jsonl")
assert os.path.exists(podcast_meta_file)

podcast_metas = []
for line in open(podcast_meta_file):
    podcast_metas.append(json.loads(line))
from suno_utils.audio import Audio
from suno_utils.tasks.hoot import (
    preload_models,
    load_model_list,
    get_word_timing_from_audio_and_lyrics,
)

_ = preload_models(
    checkpoint_filepath="/home/tony/Data/checkpoints/hoot/2025-06-23_16-49-02/last_ckpt.pt",
    tokenizer_filepath="/home/tony/Work/tony/hoot/tokenizers/v5/tokenizer_spe_bpe_v20481/tokenizer.model",
)
model_dict = load_model_list()[0]
tokenizer = model_dict["tokenizer"]
model = model_dict["model"]

from suno_utils.tasks.hoot import encode

from suno_utils.tasks.lyrics_alignment.shortest_path_aligner import (
    ShortestPathAligner,
    ShortestPathHootV6Config,
)

spa = ShortestPathAligner.from_sentencepiece(tokenizer._tokenizer, ShortestPathHootV6Config)


def align_chunk_lyrics_with_hoot(chunk_audio_path, chunk_text):
    """
    Align lyrics for a chunk using hoot alignment.

    Args:
        chunk_audio_path: Path to the chunk audio file
        chunk_text: Text content of the chunk to align

    Returns:
        dict: Alignment result with word-level timestamps
    """
    # Load audio
    audio = Audio.from_file(chunk_audio_path)

    # check cer
    hoot_result = get_word_timing_from_audio_and_lyrics(audio, chunk_text)
    cer = hoot_result[-1]["hoot_cer"]

    alignment = get_word_timing_from_audio_and_lyrics(audio, chunk_text)

    # Perform alignment
    alignment = spa.align(
        chunk_text,
        encode(audio, return_logits=True, batch_size=1),
        debug=False,
        override_enable_jumps=False,
    )

    return {
        "alignment": alignment,
        "aligned_text": "".join([w["word"] for w in alignment]),
        "original_text": chunk_text,
        "audio_path": chunk_audio_path,
        "cer": cer,
    }


def process_podcast_chunk_with_alignment(chunk):
    """
    Process a single podcast chunk and add hoot alignment.

    Args:
        chunk: Podcast chunk dict with 'audio_path' and 'text' keys

    Returns:
        dict: Enhanced chunk with alignment data
    """
    try:
        alignment_result = align_chunk_lyrics_with_hoot(chunk["audio_path"], chunk["text"])

        # Add alignment to chunk
        enhanced_chunk = chunk.copy()
        enhanced_chunk["hoot_alignment"] = alignment_result["alignment"]
        enhanced_chunk["hoot_aligned_text"] = alignment_result["aligned_text"]
        enhanced_chunk["hoot_cer"] = alignment_result["cer"]
        return enhanced_chunk
    except Exception as e:
        print(f"Error aligning chunk {chunk.get('chunk_index', 'unknown')}: {e}")
        return chunk


import json


def main():
    # Get task ID from Slurm environment variables
    task_id = int(os.environ.get("SLURM_PROCID", 0))
    total_tasks = int(os.environ.get("SLURM_NTASKS", 1))

    # Calculate which subset of podcasts this task should process
    total_podcasts = len(podcast_metas)
    podcasts_per_task = (total_podcasts + total_tasks - 1) // total_tasks
    start_idx = task_id * podcasts_per_task
    end_idx = min(start_idx + podcasts_per_task, total_podcasts)

    task_podcasts = podcast_metas[start_idx:end_idx]

    print(
        f"Task {task_id} processing podcasts {start_idx} to {end_idx - 1} ({len(task_podcasts)} podcasts)"
    )

    # Create task-specific output filename
    os.makedirs("aligned", exist_ok=True)
    output_file = f"aligned/processed_podcasts_task_{task_id:03d}.jsonl"
    with open(output_file, "a") as f:
        for record in tqdm(task_podcasts, desc=f"Task {task_id}", mininterval=60):
            try:
                record["chunks"] = [
                    process_podcast_chunk_with_alignment(chunk) for chunk in record["chunks"]
                ]
                f.write(json.dumps(record) + "\n")
            except Exception as e:
                print(f"Error processing record {record.get('id', 'unknown')}: {e}")


if __name__ == "__main__":
    main()