import os
import sys

os.environ.setdefault("HF_HOME", "/mnt/D/.cache/huggingface")
os.environ.setdefault("HF_DATASETS_CACHE", "/mnt/D/.cache/huggingface/datasets")
os.environ.setdefault("TMPDIR", "/mnt/D/tmp")

os.makedirs(os.environ["TMPDIR"], exist_ok=True)
os.makedirs(os.environ["HF_DATASETS_CACHE"], exist_ok=True)

print("=== initializing environment")

from dataclasses import dataclass
from typing import Any

import jiwer
from datasets import load_dataset, Audio
from peft import LoraConfig, TaskType, get_peft_model
import numpy as np
import torch
from tqdm import tqdm as _tqdm
from transformers import (
    Seq2SeqTrainer,
    Seq2SeqTrainingArguments,
    WhisperForConditionalGeneration,
    WhisperProcessor,
)


class WhisperLoraTrainer(Seq2SeqTrainer):
    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
        input_features = inputs.get("input_features")
        labels = inputs.get("labels")
        dtype = next(model.parameters()).dtype
        input_features = input_features.to(dtype)
        labels = labels.to(torch.long)
        outputs = model.base_model(input_features=input_features, labels=labels)
        loss = outputs.loss
        return (loss, outputs) if return_outputs else loss

    def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=None):
        if not self.args.predict_with_generate or prediction_loss_only:
            return super().prediction_step(
                model, inputs, prediction_loss_only, ignore_keys=ignore_keys
            )

        inputs = self._prepare_inputs(inputs)
        labels = inputs.get("labels")
        dtype = next(model.parameters()).dtype

        with torch.no_grad():
            generated_tokens = model.base_model.model.generate(
                input_features=inputs["input_features"].to(dtype),
                language=LANGUAGE,
                task="transcribe",
            )

        loss = None
        if labels is not None:
            with torch.no_grad():
                outputs = model.base_model(
                    input_features=inputs["input_features"].to(dtype),
                    labels=labels.to(torch.long),
                )
            loss = outputs.loss.detach().float()

        generated_tokens = generated_tokens.detach().cpu().to(torch.long)
        if labels is not None:
            labels = labels.detach().cpu().to(torch.long)

        return loss, generated_tokens, labels

print("[+] imports ok")


MODEL_ID = "openai/whisper-medium"
DATASET_ID = "fsicoli/common_voice_17_0"
LANGUAGE = "ru"
SAMPLING_RATE = 16000
OUTPUT_DIR = "./whisper-medium-ru-lora"
EVAL_SAMPLES = 1000
LORA_R = 32
LORA_ALPHA = 64
LORA_DROPOUT = 0.05


@dataclass
class DataCollatorSpeechSeq2SeqWithPadding:
    processor: Any
    decoder_start_token_id: int

    def __call__(self, features: list[dict]) -> dict:
        input_features = [{"input_features": f["input_features"]} for f in features]
        batch = self.processor.feature_extractor.pad(input_features, return_tensors="pt")

        label_features = [{"input_ids": f["labels"]} for f in features]
        labels_batch = self.processor.tokenizer.pad(label_features, return_tensors="pt")

        labels = labels_batch["input_ids"].masked_fill(
            labels_batch.attention_mask.ne(1), -100
        )

        if (labels[:, 0] == self.decoder_start_token_id).all().cpu().item():
            labels = labels[:, 1:]

        batch["labels"] = labels
        return batch


def load_processor_and_model():
    print("=== loading processor")
    try:
        processor = WhisperProcessor.from_pretrained(
            MODEL_ID, language=LANGUAGE, task="transcribe"
        )
        print("[+] processor loaded")
    except Exception as e:
        print(f"[-] processor load failed: {e}")
        sys.exit(1)

    print("=== loading model")
    try:
        model = WhisperForConditionalGeneration.from_pretrained(MODEL_ID)
        model.config.use_cache = False
        model.config.forced_decoder_ids = None
        model.config.suppress_tokens = []
        model.generation_config.language = LANGUAGE
        model.generation_config.task = "transcribe"
        model.generation_config.forced_decoder_ids = None
        model.generation_config.max_length = None
        print("[+] model loaded")
    except Exception as e:
        print(f"[-] model load failed: {e}")
        sys.exit(1)

    return processor, model


def apply_lora(model):
    print("=== applying LoRA")
    try:
        lora_config = LoraConfig(
            r=LORA_R,
            lora_alpha=LORA_ALPHA,
            lora_dropout=LORA_DROPOUT,
            target_modules=["q_proj", "k_proj", "v_proj", "out_proj"],
            task_type=TaskType.SEQ_2_SEQ_LM,
        )
        model = get_peft_model(model, lora_config)
        model.enable_input_require_grads()
        model.print_trainable_parameters()
        print("[+] LoRA applied")
    except Exception as e:
        print(f"[-] LoRA apply failed: {e}")
        sys.exit(1)
    return model


def prepare_batch(batch, processor):
    audio = batch["audio"]
    batch["input_features"] = processor.feature_extractor(
        audio["array"], sampling_rate=audio["sampling_rate"]
    ).input_features[0]
    batch["labels"] = processor.tokenizer(batch["sentence"]).input_ids
    return batch


def load_and_prepare_dataset(processor):
    print("=== loading dataset (train + validation only)")
    try:
        raw_train, raw_eval = load_dataset(
            DATASET_ID,
            LANGUAGE,
            split=["train", "validation"],
            trust_remote_code=True,
        )
        print("[+] dataset loaded")
    except Exception as e:
        print(f"[-] dataset load failed: {e}")
        sys.exit(1)

    print(f"=== slicing validation to {EVAL_SAMPLES} samples")
    raw_eval = raw_eval.select(range(EVAL_SAMPLES))
    print(f"[+] validation sliced to {len(raw_eval)} samples")

    print("=== casting audio column")
    try:
        raw_train = raw_train.cast_column("audio", Audio(sampling_rate=SAMPLING_RATE))
        raw_eval = raw_eval.cast_column("audio", Audio(sampling_rate=SAMPLING_RATE))
        print("[+] audio cast done")
    except Exception as e:
        print(f"[-] audio cast failed: {e}")
        sys.exit(1)

    fn = lambda batch: prepare_batch(batch, processor)

    print("=== mapping train split")
    try:
        train_dataset = raw_train.map(
            fn,
            remove_columns=raw_train.column_names,
            num_proc=1,
            cache_file_name=str(os.path.join(os.environ["HF_DATASETS_CACHE"], "cv17_ru_train.arrow")),
        )
        print("[+] train split mapped")
    except Exception as e:
        print(f"[-] train map failed: {e}")
        sys.exit(1)

    print("=== mapping validation split")
    try:
        eval_dataset = raw_eval.map(
            fn,
            remove_columns=raw_eval.column_names,
            num_proc=1,
            load_from_cache_file=False,
        )
        print("[+] validation split mapped")
    except Exception as e:
        print(f"[-] validation map failed: {e}")
        sys.exit(1)

    return train_dataset, eval_dataset


def build_compute_metrics(processor):
    def compute_metrics(pred):
        pred_ids = pred.predictions
        label_ids = pred.label_ids

        if isinstance(pred_ids, tuple):
            pred_ids = pred_ids[0]

        label_ids = np.array(label_ids, dtype=np.int64)
        label_ids[label_ids == -100] = processor.tokenizer.pad_token_id

        pred_str = [
            processor.tokenizer.decode(
                [int(t) for t in seq if 0 <= int(t) < processor.tokenizer.vocab_size],
                skip_special_tokens=True,
            )
            for seq in pred_ids
        ]
        label_str = processor.tokenizer.batch_decode(label_ids, skip_special_tokens=True)

        wer = 100 * jiwer.wer(label_str, pred_str)
        _tqdm.write(f"[+] eval wer: {wer:.2f}%")
        return {"wer": wer}

    return compute_metrics


def main():
    processor, model = load_processor_and_model()
    model = apply_lora(model)
    train_dataset, eval_dataset = load_and_prepare_dataset(processor)

    data_collator = DataCollatorSpeechSeq2SeqWithPadding(
        processor=processor,
        decoder_start_token_id=model.config.decoder_start_token_id,
    )
    print("[+] data collator configured")

    compute_metrics = build_compute_metrics(processor)

    print("=== configuring training arguments")
    training_args = Seq2SeqTrainingArguments(
        output_dir=OUTPUT_DIR,
        per_device_train_batch_size=2,
        gradient_accumulation_steps=1,
        learning_rate=5e-5,
        warmup_steps=500,
        max_steps=4000,
        gradient_checkpointing=True,
        predict_with_generate=True,
        per_device_eval_batch_size=2,
        generation_max_length=64,
        save_steps=1000,
        logging_steps=100,
        fp16=True,
        fp16_full_eval=False,
        eval_strategy="steps",
        eval_steps=1000,
        metric_for_best_model="wer",
        greater_is_better=False,
        load_best_model_at_end=True,
        generation_config=None,
        report_to=["tensorboard"],
    )
    print("[+] training arguments configured")

    print("=== initializing trainer")
    try:
        trainer = WhisperLoraTrainer(
            model=model,
            args=training_args,
            train_dataset=train_dataset,
            eval_dataset=eval_dataset,
            data_collator=data_collator,
            compute_metrics=compute_metrics,
            processing_class=processor.feature_extractor,
        )
        print("[+] trainer initialized")
    except Exception as e:
        print(f"[-] trainer init failed: {e}")
        sys.exit(1)

    print("=== starting training")
    trainer.train()
    print("[+] training complete")

    print("=== saving model")
    try:
        model.save_pretrained(OUTPUT_DIR)
        processor.save_pretrained(OUTPUT_DIR)
        print(f"[+] model saved to {OUTPUT_DIR}")
    except Exception as e:
        print(f"[-] save failed: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()