"""
IndicConformer Validator
========================

Uses AI4Bharat's IndicConformer model - a 600M parameter multilingual ASR.
https://huggingface.co/ai4bharat/indicconformer_stt_te_hybrid_ctc_rnnt_large

Features:
- 600M parameter Conformer architecture
- Hybrid CTC-RNNT decoding
- Fine-tuned on 22 Indian languages
- High accuracy for Indic languages
"""
import os
import time
from typing import Optional, List, Dict, Any
from pathlib import Path

from .base import (
    BaseValidator,
    ValidationResult,
    WordAlignment,
    ValidatorStatus,
    normalize_language_code
)


class IndicConformerValidator(BaseValidator):
    """
    Validator using AI4Bharat's IndicConformer model.
    
    This is a 600M parameter Conformer model trained on Indian languages.
    Can use NeMo toolkit or HuggingFace transformers.
    """
    
    name = "indic_conformer"
    description = "AI4Bharat IndicConformer - 600M Multilingual Indian ASR"
    
    # HuggingFace model IDs for language-specific models
    HF_MODELS = {
        "te": "ai4bharat/indicconformer_stt_te_hybrid_ctc_rnnt_large",
        "hi": "ai4bharat/indicconformer_stt_hi_hybrid_ctc_rnnt_large",
        "bn": "ai4bharat/indicconformer_stt_bn_hybrid_ctc_rnnt_large",
        "ta": "ai4bharat/indicconformer_stt_ta_hybrid_ctc_rnnt_large",
        "kn": "ai4bharat/indicconformer_stt_kn_hybrid_ctc_rnnt_large",
        "ml": "ai4bharat/indicconformer_stt_ml_hybrid_ctc_rnnt_large",
        "mr": "ai4bharat/indicconformer_stt_mr_hybrid_ctc_rnnt_large",
        "gu": "ai4bharat/indicconformer_stt_gu_hybrid_ctc_rnnt_large",
        "pa": "ai4bharat/indicconformer_stt_pa_hybrid_ctc_rnnt_large",
        "or": "ai4bharat/indicconformer_stt_or_hybrid_ctc_rnnt_large",
        # Multilingual fallback
        "multilingual": "ai4bharat/indic-conformer-600m-multilingual",
    }
    
    def __init__(
        self,
        enabled: bool = True,
        device: str = "auto",
        language: str = "te",
        use_nemo: bool = True,  # Try NeMo first, fallback to transformers
        **kwargs
    ):
        """
        Initialize IndicConformer validator.
        
        Args:
            enabled: Whether validator is active
            device: "cuda", "cpu", or "auto"
            language: Default language code
            use_nemo: Try to use NeMo toolkit (recommended)
        """
        super().__init__(enabled=enabled, **kwargs)
        self.device_preference = device
        self.language = normalize_language_code(language)
        self.use_nemo = use_nemo
        self.model = None
        self.device = None
        self._model_name = None
        self._model_type = None  # "nemo" or "transformers"
        self._current_language = None
        
    def setup(self, language: str = None) -> bool:
        """Load IndicConformer model."""
        try:
            import torch
            
            # Determine device
            if self.device_preference == "auto":
                self.device = "cuda" if torch.cuda.is_available() else "cpu"
            else:
                self.device = self.device_preference
                
            print(f"[{self.name}] Loading model on {self.device}...")
            
            lang_code = normalize_language_code(language or self.language)
            self._current_language = lang_code
            
            # Try NeMo first
            if self.use_nemo:
                try:
                    return self._setup_nemo(lang_code)
                except ImportError:
                    print(f"[{self.name}] NeMo not available, trying transformers...")
                except Exception as e:
                    print(f"[{self.name}] NeMo setup failed: {e}")
            
            # Fallback to transformers
            return self._setup_transformers(lang_code)
            
        except Exception as e:
            print(f"[{self.name}] Setup error: {e}")
            import traceback
            traceback.print_exc()
            return False
    
    def _setup_nemo(self, language: str) -> bool:
        """Setup using NeMo toolkit."""
        try:
            import nemo.collections.asr as nemo_asr
        except ImportError:
            print(f"[{self.name}] NeMo not installed. Install with: pip install nemo_toolkit[asr]")
            raise ImportError("NeMo toolkit required for IndicConformer")
        
        # Get model name for language
        self._model_name = self.HF_MODELS.get(
            language,
            self.HF_MODELS["multilingual"]
        )
        
        print(f"[{self.name}] Loading NeMo model: {self._model_name}")
        
        # Load from HuggingFace
        self.model = nemo_asr.models.ASRModel.from_pretrained(
            model_name=self._model_name
        )
        
        if self.device == "cuda":
            self.model = self.model.cuda()
        
        self.model.eval()
        self._model_type = "nemo"
        
        print(f"[{self.name}] NeMo model loaded successfully")
        return True
    
    def _setup_transformers(self, language: str) -> bool:
        """Setup using HuggingFace transformers."""
        # Note: IndicConformer models are NeMo models and require NeMo toolkit
        # They cannot be loaded directly via transformers
        print(f"[{self.name}] IndicConformer models require NeMo toolkit")
        print(f"[{self.name}] Install with: pip install nemo_toolkit[asr]")
        print(f"[{self.name}] Then run: python -c \"import nemo.collections.asr as nemo_asr; nemo_asr.models.ASRModel.from_pretrained('ai4bharat/indicconformer_stt_te_hybrid_ctc_rnnt_large')\"")
        return False
    
    def _load_audio(self, audio_path: str, target_sr: int = 16000):
        """Load and resample audio."""
        import soundfile as sf
        import torch
        import numpy as np
        
        data, sample_rate = sf.read(audio_path)
        waveform = torch.from_numpy(data).float()
        
        if len(waveform.shape) == 1:
            waveform = waveform.unsqueeze(0)
        
        if waveform.shape[0] > 1:
            waveform = waveform.mean(dim=0, keepdim=True)
        
        if sample_rate != target_sr:
            import torchaudio
            resampler = torchaudio.transforms.Resample(sample_rate, target_sr)
            waveform = resampler(waveform)
            
        return waveform.squeeze(), target_sr
    
    def _transcribe_nemo(self, audio_path: str) -> Dict[str, Any]:
        """Transcribe using NeMo model."""
        # NeMo can transcribe directly from file path
        transcription = self.model.transcribe([audio_path])[0]
        
        # Get duration
        import soundfile as sf
        data, sr = sf.read(audio_path)
        duration = len(data) / sr
        
        # Generate simple word alignments
        words = transcription.split()
        alignments = []
        
        for i, word in enumerate(words):
            alignments.append(WordAlignment(
                word=word,
                start_time=i * duration / max(len(words), 1),
                end_time=(i + 1) * duration / max(len(words), 1),
                confidence=None
            ))
        
        return {
            "transcription": transcription,
            "alignments": alignments,
            "duration": duration
        }
    
    def _transcribe_transformers(self, audio_path: str) -> Dict[str, Any]:
        """Transcribe using transformers model."""
        import torch
        
        waveform, sr = self._load_audio(audio_path)
        duration = len(waveform) / sr
        
        inputs = self.processor(
            waveform.numpy(),
            sampling_rate=sr,
            return_tensors="pt",
            padding=True
        )
        
        input_features = inputs.input_values if hasattr(inputs, 'input_values') else inputs.input_features
        input_features = input_features.to(self.device)
        
        with torch.no_grad():
            outputs = self.model(input_features)
            
            if hasattr(outputs, 'logits'):
                predicted_ids = torch.argmax(outputs.logits, dim=-1)
                transcription = self.processor.batch_decode(predicted_ids)[0]
            else:
                # Seq2Seq model
                generated_ids = self.model.generate(input_features)
                transcription = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
        
        # Generate simple word alignments
        words = transcription.split()
        alignments = []
        
        for i, word in enumerate(words):
            alignments.append(WordAlignment(
                word=word,
                start_time=i * duration / max(len(words), 1),
                end_time=(i + 1) * duration / max(len(words), 1),
                confidence=None
            ))
        
        return {
            "transcription": transcription,
            "alignments": alignments,
            "duration": duration
        }
    
    def validate(
        self,
        audio_path: str,
        reference_text: Optional[str] = None,
        language: str = "te"
    ) -> ValidationResult:
        """
        Transcribe audio using IndicConformer.
        
        Args:
            audio_path: Path to audio file
            reference_text: Optional reference (not used)
            language: Language code
            
        Returns:
            ValidationResult with ASR output
        """
        if not self.enabled:
            return ValidationResult(
                validator_name=self.name,
                audio_path=audio_path,
                success=False,
                error_message="Validator disabled"
            )
        
        start_time = time.time()
        lang_code = normalize_language_code(language)
        
        # Check if we need to reload model for different language
        if self._current_language != lang_code or self.model is None:
            if not self.setup(lang_code):
                return ValidationResult(
                    validator_name=self.name,
                    audio_path=audio_path,
                    success=False,
                    error_message="Model not loaded"
                )
        
        try:
            # Transcribe based on model type
            if self._model_type == "nemo":
                result = self._transcribe_nemo(audio_path)
            else:
                result = self._transcribe_transformers(audio_path)
            
            processing_time = time.time() - start_time
            
            return ValidationResult(
                validator_name=self.name,
                audio_path=audio_path,
                success=True,
                transcription=result["transcription"],
                word_alignments=result["alignments"],
                overall_confidence=None,  # Conformer doesn't provide confidence easily
                processing_time_sec=processing_time,
                audio_duration_sec=result.get("duration"),
                raw_output={
                    "language": lang_code,
                    "model": self._model_name,
                    "model_type": self._model_type
                }
            )
            
        except Exception as e:
            import traceback
            traceback.print_exc()
            return ValidationResult(
                validator_name=self.name,
                audio_path=audio_path,
                success=False,
                error_message=str(e),
                processing_time_sec=time.time() - start_time
            )
    
    def cleanup(self):
        """Release resources."""
        if self.model is not None:
            del self.model
            self.model = None
            
        import torch
        if torch.cuda.is_available():
            torch.cuda.empty_cache()