import argparse
from pathlib import Path

import sacrebleu
import torch
import tqdm
import yaml
from torch.utils.data import DataLoader
from transformers import AutoModel, AutoTokenizer

from src.model import SLTModel
from src.train import SignLanguageDataset


def collate_fn(batch):
    # Minimal collator for inference loop
    signs = [item["sign"] for item in batch]
    texts = [item["text"] for item in batch]
    return signs, texts


# --- Generation Logic ---
def generate_batch(
    model, tokenizer, sign_features_list, max_new_tokens=60, device="cpu"
):
    """
    Greedy decoding for a batch of sign sequences.
    Note: Ideally, you should pad signs and run in batch mode.
    For simplicity and safety with variable length audio/signs, we loop here
    or you can implement padding like the training collator.

    Here we implement single-sample processing loop for safety to avoid padding issues during inference
    if lengths vary wildly, but it's slower.
    """
    model.eval()
    translations = []

    for sign in sign_features_list:
        # Prepare Input: [1, seq_len, dim]
        input_values = sign.unsqueeze(0).to(device)
        attention_mask = torch.ones(
            input_values.shape[0], input_values.shape[1], device=device
        )

        # Start with BOS
        decoder_input_ids = torch.tensor([[tokenizer.cls_token_id]], device=device)

        # Decode Loop
        for _ in range(max_new_tokens):
            with torch.no_grad():
                outputs = model(
                    input_values=input_values,
                    attention_mask=attention_mask,
                    decoder_input_ids=decoder_input_ids,
                )

            next_token_logits = outputs.logits[:, -1, :]
            next_token_id = torch.argmax(next_token_logits, dim=-1).unsqueeze(-1)

            decoder_input_ids = torch.cat([decoder_input_ids, next_token_id], dim=1)

            if next_token_id.item() == tokenizer.sep_token_id:
                break

        # Decode
        pred_text = tokenizer.decode(decoder_input_ids[0], skip_special_tokens=True)
        translations.append(pred_text)

    return translations


def parse_args():
    parser = argparse.ArgumentParser(description="SLT Inference Script")
    parser.add_argument(
        "--config",
        type=str,
        default="configs/config.yaml",
        help="Path to the config file",
    )
    parser.add_argument(
        "--model_path",
        type=str,
        help="Path to the trained model",
    )
    parser.add_argument(
        "--output_file",
        type=str,
        default="predictions.txt",
        help="File to save the predictions",
    )
    parser.add_argument(
        "--use_hub",
        action="store_true",
        help="Whether to load model from HuggingFace Hub",
    )
    return parser.parse_args()


def main():
    # load config file
    args = parse_args()
    with open(Path(args.config)) as f:
        try:
            config_file = yaml.safe_load(f)
        except yaml.YAMLError as exc:
            raise exc

    # Settings
    device = "cuda" if torch.cuda.is_available() else "cpu"
    batch_size = 1

    # 1. Load Model & Tokenizer
    if not args.use_hub:
        print(f"Loading model from {args.model_path}...")
        tokenizer = AutoTokenizer.from_pretrained(args.model_path)
        model = SLTModel.from_pretrained(args.model_path).to(device)
    else:
        print(f"Loading model from HuggingFace Hub: {args.model_path}...")
        tokenizer = AutoTokenizer.from_pretrained(
            Path(args.model_path), trust_remote_code=True
        )
        model = AutoModel.from_pretrained(
            Path(args.model_path), trust_remote_code=True
        ).to(device)

    # 2. Load Dataset
    test_dataset = SignLanguageDataset(config_file["dataset"]["pickle_path"]["test"])
    test_loader = DataLoader(
        test_dataset, batch_size=batch_size, shuffle=False, collate_fn=collate_fn
    )

    # 3. Inference Loop
    all_predictions = []
    all_references = []

    print(f"Starting inference on {len(test_dataset)} samples...")

    with open(Path(args.output_file), "w") as f_out:
        for batch_signs, batch_texts in tqdm.tqdm(test_loader):
            # batch_signs is a list of tensors
            preds = generate_batch(model, tokenizer, batch_signs, device=device)

            for pred, ref in zip(preds, batch_texts):
                all_predictions.append(pred)
                all_references.append(ref)

                f_out.write(f"REF: {ref}\n")
                f_out.write(f"HYP: {pred}\n")
                f_out.write("-" * 20 + "\n")

    print(f"Inference done. Results saved to {args.output_file}")

    # 4. Calculate Metric
    bleu = sacrebleu.corpus_bleu(all_predictions, [all_references])
    print(f"BLEU Score: {bleu.score}")


if __name__ == "__main__":
    main()