"""
Translator - Main public API for BISINDO Trans.

This module provides the unified Translator class that combines
preprocessing, models, decoding, and postprocessing.
"""

import os
from typing import Literal, Optional

from bisindotrans.preprocessing import normalisasi_glosa, naturalisasi_aturan
from bisindotrans.utils import bersihkan_output


class Translator:
    """
    Penerjemah Bahasa Isyarat Indonesia ke Bahasa Indonesia.
    
    Supports two model types:
        - "ngram": Statistical N-gram model (fast, offline)
        - "neural": Neural Language Model GPT-2 (requires torch)
    
    Example:
        >>> from bisindotrans import Translator
        >>> t = Translator(model_type="ngram")
        >>> t.translate("SAYA MAKAN NASI")
        'Saya makan nasi.'
    """
    
    def __init__(
        self,
        model_type: Literal["ngram", "neural"] = "ngram",
        model_path: Optional[str] = None,
        auto_load: bool = True,
    ):
        """
        Initialize the Translator.
        
        Args:
            model_type: "ngram" or "neural"
            model_path: Path to saved model (for ngram: .pkl file)
            auto_load: If True, automatically load/train model
        """
        self.model_type = model_type
        self.model = None
        self.decoder = None
        self._is_ready = False
        
        if auto_load:
            self._initialize_model(model_path)
    
    def _initialize_model(self, model_path: Optional[str] = None):
        """Load or initialize the language model."""
        
        if self.model_type == "ngram":
            self._init_ngram(model_path)
        elif self.model_type == "neural":
            self._init_neural()
        else:
            raise ValueError(f"Unknown model_type: {self.model_type}")
        
        self._is_ready = True
    
    def _init_ngram(self, model_path: Optional[str] = None):
        """Initialize N-gram model with persistence support."""
        from bisindotrans.models.ngram import ModelNgram
        from bisindotrans.decoding.strategies import Decoding
        
        self.model = ModelNgram()
        
        # Default model path
        if model_path is None:
            model_path = "model/ngram_trained.pkl"
        
        # Try to load saved model
        if os.path.exists(model_path):
            print(f"[Translator] Loading saved model: {model_path}")
            self.model.muat_model(model_path)
        else:
            print("[Translator] No saved model found. Training from corpus...")
            corpus_files = [
                "data/corpus_indonesia.txt",
                "data/corpus_indonesia_2.txt",
                "data/corpus_indonesia_3.txt",
            ]
            for path in corpus_files:
                if os.path.exists(path):
                    self.model.latih(path)
            
            # Save for next time
            os.makedirs(os.path.dirname(model_path), exist_ok=True)
            self.model.simpan_model(model_path)
        
        self.decoder = Decoding(self.model)
        print(f"[Translator] N-gram model ready. Vocab: {len(self.model.unigram_counts):,}")
    
    def _init_neural(self):
        """Initialize Neural LM (GPT-2)."""
        try:
            from bisindotrans.models.neural import ModelNeural
            from bisindotrans.decoding.strategies import NeuralDecoding
        except ImportError as e:
            raise ImportError(
                "Neural model requires 'torch' and 'transformers'. "
                "Install with: pip install bisindo-trans[neural]"
            ) from e
        
        self.model = ModelNeural()
        self.decoder = NeuralDecoding(self.model)
        print("[Translator] Neural model ready.")
    
    def translate(
        self,
        glosa_input: str,
        method: Literal["beam", "nucleus"] = "beam",
        **kwargs
    ) -> str:
        """
        Translate glosa input to natural Indonesian text.
        
        Args:
            glosa_input: Space-separated glosa string (e.g., "SAYA MAKAN NASI")
            method: Decoding method - "beam" (deterministic) or "nucleus" (stochastic)
            **kwargs: Additional arguments for decoder (beam_width, top_p, etc.)
            
        Returns:
            Translated Indonesian text
            
        Example:
            >>> t.translate("SAYA MAKAN NASI", method="beam")
            'Saya makan nasi.'
        """
        if not self._is_ready:
            raise RuntimeError("Model not initialized. Call _initialize_model() first.")
        
        # Parse input
        if isinstance(glosa_input, str):
            input_list = glosa_input.upper().split()
        else:
            input_list = list(glosa_input)
        
        # Preprocessing pipeline
        normalized = normalisasi_glosa(input_list)
        naturalized = naturalisasi_aturan(normalized)
        
        # Decoding
        if method == "beam":
            if self.model_type == "ngram":
                raw_output = self.decoder.beam_search(naturalized, beam_width=kwargs.get("beam_width", 3))
            else:
                raw_output = self.decoder.beam_search(naturalized, beam_width=kwargs.get("beam_width", 5))
        elif method == "nucleus":
            if self.model_type == "ngram":
                raw_output = self.decoder.nucleus_sampling(naturalized, top_p=kwargs.get("top_p", 0.5))
            else:
                raw_output = self.decoder.nucleus_sampling(
                    naturalized, 
                    top_p=kwargs.get("top_p", 0.7),
                    temperature=kwargs.get("temperature", 0.6)
                )
        else:
            raise ValueError(f"Unknown method: {method}. Use 'beam' or 'nucleus'.")
        
        # Postprocessing
        final_output = bersihkan_output(raw_output)
        
        return final_output
    
    def translate_batch(self, glosa_list: list, method: str = "beam", **kwargs) -> list:
        """
        Translate a batch of glosa inputs.
        
        Args:
            glosa_list: List of glosa strings
            method: Decoding method
            **kwargs: Additional decoder arguments
            
        Returns:
            List of translated strings
        """
        return [self.translate(g, method=method, **kwargs) for g in glosa_list]
    
    @property
    def is_ready(self) -> bool:
        """Check if translator is initialized and ready."""
        return self._is_ready
    
    def __repr__(self) -> str:
        status = "ready" if self._is_ready else "not initialized"
        return f"Translator(model_type='{self.model_type}', status='{status}')"


# --- TEST ---
if __name__ == "__main__":
    print("=== Testing Translator ===\n")
    
    # Test with N-gram
    t = Translator(model_type="ngram")
    
    test_cases = [
        "SAYA MAKAN NASI",
        "SELAMAT PAGI",
        "TERIMA KASIH IBU",
        "SIAPA NAMA KAMU",
    ]
    
    print("\n--- Results ---")
    for glosa in test_cases:
        result = t.translate(glosa, method="beam")
        print(f"{glosa} -> {result}")
