"""
Test script: Download and verify Whisper-small (ASR) and Qwen2.5-3B-Instruct (Q&A).
Run this to confirm models load correctly before wiring into app.py.

Usage:
    python test_modules/test_whisper_qwen.py
"""
import sys
import time

print("=" * 60)
print("Step 1: Testing Whisper-small (ASR)")
print("=" * 60)

try:
    import torch
    from transformers import pipeline

    start = time.time()
    print("Loading whisper-small pipeline...")
    asr_pipe = pipeline(
        "automatic-speech-recognition",
        model="openai/whisper-small",
        device="cuda" if torch.cuda.is_available() else "cpu",
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    )
    elapsed = time.time() - start
    print(f"[OK] Whisper-small loaded in {elapsed:.1f}s")
    print(f"  Device: {'cuda' if torch.cuda.is_available() else 'cpu'}")

    # Test with a short synthetic audio array
    import numpy as np
    dummy_audio = np.zeros(16000, dtype=np.float32)  # 1 second of silence
    result = asr_pipe({"raw": dummy_audio, "sampling_rate": 16000})
    print(f"[OK] Whisper inference test passed (result: '{result['text'].strip()}')")

except Exception as e:
    print(f"[FAIL] Whisper-small failed: {e}")
    sys.exit(1)

print()
print("=" * 60)
print("Step 2: Testing Qwen2.5-3B-Instruct (Q&A)")
print("=" * 60)

try:
    from transformers import AutoTokenizer, AutoModelForCausalLM

    start = time.time()
    model_id = "Qwen/Qwen2.5-3B-Instruct"
    print(f"Loading {model_id}...")

    tokenizer = AutoTokenizer.from_pretrained(model_id)

    # Use float16 on GPU, float32 on CPU. 4-bit quantization used on HF Spaces (Linux).
    if torch.cuda.is_available():
        model = AutoModelForCausalLM.from_pretrained(
            model_id,
            torch_dtype=torch.float16,
            device_map="auto",
        )
    else:
        print("  (No GPU -- loading in float32 on CPU, will be slow)")
        model = AutoModelForCausalLM.from_pretrained(
            model_id,
            torch_dtype=torch.float32,
            device_map="cpu",
        )

    elapsed = time.time() - start
    print(f"[OK] Qwen2.5-3B-Instruct loaded in {elapsed:.1f}s")

    # Test inference with a story Q&A prompt
    story_context = "Peter Rabbit squeezed under the gate into Mr. McGregor's garden. He ate some lettuces and French beans."
    question = "What did Peter Rabbit eat?"

    messages = [
        {"role": "system", "content": "You are a friendly storyteller answering a child's question about a story. Answer in 1-2 short sentences using only information from the story context provided."},
        {"role": "user", "content": f"Story context: {story_context}\n\nChild's question: {question}"}
    ]

    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(text, return_tensors="pt").to(model.device)

    start = time.time()
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=80,
            temperature=0.7,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id,
        )
    answer_tokens = outputs[0][inputs["input_ids"].shape[1]:]
    answer = tokenizer.decode(answer_tokens, skip_special_tokens=True)
    elapsed = time.time() - start

    print(f"[OK] Qwen inference test passed in {elapsed:.1f}s")
    print(f"  Q: {question}")
    print(f"  A: {answer}")

except Exception as e:
    print(f"[FAIL] Qwen2.5-3B-Instruct failed: {e}")
    import traceback
    traceback.print_exc()
    sys.exit(1)

print()
print("=" * 60)
print("[OK] ALL MODELS VERIFIED -- ready for integration")
print("=" * 60)