# evaluate.py - Handles evaluation and comparing tasks

import os
import glob
import logging
import traceback
import tempfile
import shutil
from difflib import SequenceMatcher
import torch
import torchaudio
from pydub import AudioSegment
from flask import jsonify
from werkzeug.utils import secure_filename
from concurrent.futures import ThreadPoolExecutor

# Import necessary functions from translator.py
from translator import get_asr_model, get_asr_processor, LANGUAGE_CODES

# Configure logging
logger = logging.getLogger("speech_api")

def calculate_similarity(text1, text2):
    """Calculate text similarity percentage."""
    def clean_text(text):
        return text.lower()

    clean1 = clean_text(text1)
    clean2 = clean_text(text2)

    matcher = SequenceMatcher(None, clean1, clean2)
    return matcher.ratio() * 100

def setup_reference_patterns(reference_dir, sample_rate=16000):
    """Create standard reference pattern directories and dummy files if needed"""
    reference_patterns = [
        "mayap_a_abak", "mayap_a_ugtu", "mayap_a_gatpanapun", "mayap_a_bengi", 
        "komusta_ka", "malaus_ko_pu", "malaus_kayu", "agaganaka_da_ka", 
        "pagdulapan_da_ka", "kaluguran_da_ka", "dakal_a_salamat", "panapaya_mu_ku",
        "wa", "ali", "tuknang", "lagwa", "galo", "buri_ke_ini", "tara_na",  
        "nokarin_ka_ibat", "nokarin_ka_munta", "atiu_na_ku", "nanung_panayan_mu",  
        "mako_na_ka", "muli_ta_na", "nanu_ing_pengan_mu", "mekeni", "mengan_na_ka",   
        "munta_ka_karin", "magkanu_ini", "mimingat_ka", "mangan_ta_na", "lakwan_da_ka",  
        "nanu_maliari_kung_daptan_keka", "pilan_na_ka_banwa", "saliwan_ke_ini",  
        "makananu_munta_king"
    ]
    
    created_dirs = 0
    created_files = 0
    
    for pattern in reference_patterns:
        pattern_dir = os.path.join(reference_dir, pattern)
        if not os.path.exists(pattern_dir):
            try:
                os.makedirs(pattern_dir, exist_ok=True)
                logger.info(f"📁 Created reference pattern directory: {pattern_dir}")
                created_dirs += 1
            except Exception as e:
                logger.error(f"❌ Failed to create reference pattern directory {pattern_dir}: {str(e)}")
                continue
        
        # Check if directory has any WAV files, add a dummy if not
        wav_files = glob.glob(os.path.join(pattern_dir, "*.wav"))
        if not wav_files:
            try:
                dummy_path = os.path.join(pattern_dir, "dummy_reference.wav")
                # Create a 1-second silent audio file - not completely silent to avoid transcription issues
                # Adding a small amount of noise helps ASR models detect something
                silent = AudioSegment.silent(duration=1000, frame_rate=sample_rate)
                # Add a tiny bit of noise
                for i in range(50, 950, 300):
                    silent = silent.overlay(AudioSegment.silent(duration=50, frame_rate=sample_rate) + 3, position=i)
                silent.export(dummy_path, format="wav")
                logger.info(f"📄 Created dummy reference file: {dummy_path}")
                created_files += 1
            except Exception as e:
                logger.error(f"❌ Failed to create dummy file in {pattern_dir}: {str(e)}")
    
    return created_dirs, created_files

def search_reference_directories():
    """Search for possible reference directories in various locations"""
    possible_locations = [
        "./reference_audios",
        "../reference_audios",
        "/app/reference_audios",
        "/tmp/reference_audios",
        os.path.join(os.path.dirname(os.path.abspath(__file__)), "reference_audios")
    ]
    
    found_dirs = []
    for location in possible_locations:
        if os.path.exists(location) and os.path.isdir(location):
            access_info = {
                "readable": os.access(location, os.R_OK),
                "writable": os.access(location, os.W_OK),
                "executable": os.access(location, os.X_OK)
            }
            
            # Count pattern directories
            pattern_dirs = [d for d in os.listdir(location) 
                          if os.path.isdir(os.path.join(location, d))]
            
            # Count total wav files
            wav_count = 0
            for pattern in pattern_dirs:
                pattern_path = os.path.join(location, pattern)
                wav_count += len(glob.glob(os.path.join(pattern_path, "*.wav")))
            
            found_dirs.append({
                "path": location,
                "access": access_info,
                "pattern_dirs": len(pattern_dirs),
                "wav_files": wav_count
            })
    
    return found_dirs

def init_reference_audio(reference_dir, output_dir):
    """Initialize reference audio directories and return the working directory path"""
    try:
        # Create the output directory first
        os.makedirs(output_dir, exist_ok=True)
        logger.info(f"📁 Created output directory: {output_dir}")

        # Search for existing reference directories
        found_dirs = search_reference_directories()
        for directory in found_dirs:
            logger.info(f"🔍 Found reference directory: {directory['path']} "
                      f"(patterns: {directory['pattern_dirs']}, wav files: {directory['wav_files']})")
        
        # First, try to use the provided reference_dir
        working_dir = reference_dir
        
        # Check if reference_dir is accessible and writable
        if not os.path.exists(reference_dir) or not os.access(reference_dir, os.W_OK):
            logger.warning(f"⚠️ Provided reference directory {reference_dir} is not writable")
            
            # Try to use a found directory that has patterns and is writable
            for directory in found_dirs:
                if directory['access']['writable'] and directory['pattern_dirs'] > 0:
                    working_dir = directory['path']
                    logger.info(f"✅ Using found reference directory: {working_dir}")
                    break
            else:
                # If no suitable directory found, create one in /tmp
                working_dir = os.path.join('/tmp', 'reference_audios')
                logger.warning(f"⚠️ Using fallback reference directory in /tmp: {working_dir}")
        
        # Ensure the working directory exists
        os.makedirs(working_dir, exist_ok=True)
        logger.info(f"📁 Using reference directory: {working_dir}")
        
        # Set up reference pattern directories with dummy files if needed
        dirs_created, files_created = setup_reference_patterns(working_dir)
        logger.info(f"📊 Created {dirs_created} directories and {files_created} dummy files")
        
        # Try to copy reference files from other found directories to working directory if needed
        if files_created > 0 and len(found_dirs) > 1:
            # Try to find a directory with existing WAV files
            for directory in found_dirs:
                if directory['path'] != working_dir and directory['wav_files'] > 0:
                    try:
                        source_dir = directory['path']
                        logger.info(f"🔄 Copying reference files from {source_dir} to {working_dir}")
                        
                        # Copy pattern directories that have WAV files
                        for item in os.listdir(source_dir):
                            src_path = os.path.join(source_dir, item)
                            if os.path.isdir(src_path) and glob.glob(os.path.join(src_path, "*.wav")):
                                dst_path = os.path.join(working_dir, item)
                                
                                # Copy each WAV file individually
                                for wav_file in glob.glob(os.path.join(src_path, "*.wav")):
                                    wav_name = os.path.basename(wav_file)
                                    dst_file = os.path.join(dst_path, wav_name)
                                    if not os.path.exists(dst_file):
                                        shutil.copy2(wav_file, dst_file)
                                        logger.info(f"📄 Copied {wav_name} to {dst_path}")
                        
                        break
                    except Exception as e:
                        logger.warning(f"⚠️ Failed to copy reference files: {str(e)}")
        
        # Log the final contents
        pattern_dirs = [d for d in os.listdir(working_dir)
                       if os.path.isdir(os.path.join(working_dir, d))]
        logger.info(f"📊 Final reference directory has {len(pattern_dirs)} pattern directories")
        
        total_wav_files = 0
        for pattern in pattern_dirs:
            pattern_path = os.path.join(working_dir, pattern)
            wav_files = glob.glob(os.path.join(pattern_path, "*.wav"))
            total_wav_files += len(wav_files)
            logger.info(f"  - {pattern}: {len(wav_files)} WAV files")
        
        logger.info(f"📊 Total reference WAV files: {total_wav_files}")
        
        return working_dir
            
    except Exception as e:
        logger.error(f"❌ Failed to set up reference audio directory: {str(e)}")
        logger.debug(f"Stack trace: {traceback.format_exc()}")
        
        # As a last resort, try to use /tmp
        fallback_dir = os.path.join('/tmp', 'reference_audios')
        try:
            os.makedirs(fallback_dir, exist_ok=True)
            setup_reference_patterns(fallback_dir)
            logger.warning(f"⚠️ Using emergency fallback directory: {fallback_dir}")
            return fallback_dir
        except:
            logger.critical("💥 CRITICAL: Failed to create even a fallback directory")
            return reference_dir

def handle_upload_reference(request, reference_dir, sample_rate):
    """Handle upload of reference audio files"""
    try:
        if "audio" not in request.files:
            logger.warning("⚠️ Reference upload missing audio file")
            return jsonify({"error": "No audio file uploaded"}), 400

        reference_word = request.form.get("reference_word", "").strip()
        if not reference_word:
            logger.warning("⚠️ Reference upload missing reference word")
            return jsonify({"error": "No reference word provided"}), 400

        # Validate reference word
        reference_patterns = [
            "mayap_a_abak", "mayap_a_ugtu", "mayap_a_gatpanapun", "mayap_a_bengi",
            "komusta_ka", "malaus_ko_pu", "malaus_kayu", "agaganaka_da_ka",
            "pagdulapan_da_ka", "kaluguran_da_ka", "dakal_a_salamat", "panapaya_mu_ku",
            "wa", "ali", "tuknang", "lagwa", "galo", "buri_ke_ini", "tara_na",  
            "nokarin_ka_ibat", "nokarin_ka_munta", "atiu_na_ku", "nanung_panayan_mu",  
            "mako_na_ka", "muli_ta_na", "nanu_ing_pengan_mu", "mekeni", "mengan_na_ka",   
            "munta_ka_karin", "magkanu_ini", "mimingat_ka", "mangan_ta_na", "lakwan_da_ka",  
            "nanu_maliari_kung_daptan_keka", "pilan_na_ka_banwa", "saliwan_ke_ini",  
            "makananu_munta_king"
        ]

        if reference_word not in reference_patterns:
            logger.warning(f"⚠️ Invalid reference word: {reference_word}")
            return jsonify({"error": f"Invalid reference word. Available: {reference_patterns}"}), 400

        # Make sure we have a writable reference directory
        if not os.path.exists(reference_dir):
            reference_dir = os.path.join('/tmp', 'reference_audios')
            os.makedirs(reference_dir, exist_ok=True)
            logger.warning(f"⚠️ Using alternate reference directory for upload: {reference_dir}")

        # Create directory for reference pattern if it doesn't exist
        pattern_dir = os.path.join(reference_dir, reference_word)
        os.makedirs(pattern_dir, exist_ok=True)

        # Save the reference audio file
        audio_file = request.files["audio"]
        filename = secure_filename(audio_file.filename)
        
        # Ensure filename has .wav extension
        if not filename.lower().endswith('.wav'):
            base_name = os.path.splitext(filename)[0]
            filename = f"{base_name}.wav"
            
        file_path = os.path.join(pattern_dir, filename)
        
        # Create a temporary file first, then convert to WAV
        with tempfile.NamedTemporaryFile(delete=False) as temp_file:
            audio_file.save(temp_file.name)
            temp_path = temp_file.name
        
        try:
            # Process the audio file
            audio = AudioSegment.from_file(temp_path)
            audio = audio.set_frame_rate(sample_rate).set_channels(1)
            audio.export(file_path, format="wav")
            logger.info(f"✅ Reference audio saved successfully for {reference_word}: {file_path}")
            
            # Clean up temp file
            try:
                os.unlink(temp_path)
            except:
                pass
        except Exception as e:
            logger.error(f"❌ Reference audio processing failed: {str(e)}")
            return jsonify({"error": f"Audio processing failed: {str(e)}"}), 500

        # Count how many references we have now
        references = glob.glob(os.path.join(pattern_dir, "*.wav"))
        return jsonify({
            "message": "Reference audio uploaded successfully",
            "reference_word": reference_word,
            "file": filename,
            "total_references": len(references)
        })

    except Exception as e:
        logger.error(f"❌ Unhandled exception in reference upload: {str(e)}")
        logger.debug(f"Stack trace: {traceback.format_exc()}")
        return jsonify({"error": f"Internal server error: {str(e)}"}), 500

def handle_evaluation_request(request, reference_dir, output_dir, sample_rate):
    """Handle pronunciation evaluation requests"""
    request_id = f"req-{id(request)}"  # Create unique ID for this request
    logger.info(f"[{request_id}] 🆕 Starting new pronunciation evaluation request")
    
    temp_dir = None
    
    # Get the ASR model and processor using the getter functions
    asr_model = get_asr_model()
    asr_processor = get_asr_processor()
    
    if asr_model is None or asr_processor is None:
        logger.error(f"[{request_id}] ❌ Evaluation endpoint called but ASR models aren't loaded")
        return jsonify({"error": "ASR model not available"}), 503

    try:
        if "audio" not in request.files:
            logger.warning(f"[{request_id}] ⚠️ Evaluation request missing audio file")
            return jsonify({"error": "No audio file uploaded"}), 400

        audio_file = request.files["audio"]
        reference_locator = request.form.get("reference_locator", "").strip()
        language = request.form.get("language", "kapampangan").lower()

        # Validate reference locator
        if not reference_locator:
            logger.warning(f"[{request_id}] ⚠️ No reference locator provided")
            return jsonify({"error": "Reference locator is required"}), 400

        # Construct full reference directory path
        reference_dir_path = os.path.join(reference_dir, reference_locator)
        logger.info(f"[{request_id}] 📁 Reference directory path: {reference_dir_path}")

        # Make sure the reference directory exists
        if not os.path.exists(reference_dir_path):
            try:
                os.makedirs(reference_dir_path, exist_ok=True)
                logger.warning(f"[{request_id}] ⚠️ Created missing reference directory: {reference_dir_path}")
            except Exception as e:
                logger.error(f"[{request_id}] ❌ Failed to create reference directory: {str(e)}")
                return jsonify({"error": f"Reference audio directory not found: {reference_locator}"}), 404

        # Check for reference files
        reference_files = glob.glob(os.path.join(reference_dir_path, "*.wav"))
        logger.info(f"[{request_id}] 📁 Found {len(reference_files)} reference files")

        # If no reference files exist, create a dummy reference file
        if not reference_files:
            logger.warning(f"[{request_id}] ⚠️ No reference audio files found in {reference_dir_path}")
    
            # Create a dummy reference file
            try:
                dummy_file_path = os.path.join(reference_dir_path, "dummy_reference.wav")
                logger.info(f"[{request_id}] 🔄 Creating dummy reference file: {dummy_file_path}")
        
                # Create a 1-second audio file with a slight sound
                silent_audio = AudioSegment.silent(duration=1000, frame_rate=sample_rate)
                # Add a tiny bit of noise to help ASR
                for i in range(50, 950, 300):
                    silent_audio = silent_audio.overlay(AudioSegment.silent(duration=50, frame_rate=sample_rate) + 3, position=i)
                silent_audio.export(dummy_file_path, format="wav")
        
                # Add it to the list of reference files
                reference_files = [dummy_file_path]
                logger.info(f"[{request_id}] ✅ Created dummy reference file for testing")
            except Exception as e:
                logger.error(f"[{request_id}] ❌ Failed to create dummy reference: {str(e)}")
                return jsonify({"error": f"No reference audio found for {reference_locator}"}), 404

        lang_code = LANGUAGE_CODES.get(language, language)
        logger.info(f"[{request_id}] 🔄 Evaluating pronunciation for reference: {reference_locator} with language code: {lang_code}")

        # Create a request-specific temp directory to avoid conflicts
        temp_dir = os.path.join(output_dir, f"temp_{request_id}")
        os.makedirs(temp_dir, exist_ok=True)

        # Process user audio
        user_audio_path = os.path.join(temp_dir, "user_audio_input.wav")
        with open(user_audio_path, 'wb') as f:
            f.write(audio_file.read())

        try:
            logger.info(f"[{request_id}] 🔄 Processing user audio file")
            audio = AudioSegment.from_file(user_audio_path)
            audio = audio.set_frame_rate(sample_rate).set_channels(1)

            processed_path = os.path.join(temp_dir, "processed_user_audio.wav")
            audio.export(processed_path, format="wav")

            user_waveform, sr = torchaudio.load(processed_path)
            user_waveform = user_waveform.squeeze().numpy()
            logger.info(f"[{request_id}] ✅ User audio processed: {sr}Hz, length: {len(user_waveform)} samples")

            user_audio_path = processed_path
        except Exception as e:
            logger.error(f"[{request_id}] ❌ Audio processing failed: {str(e)}")
            return jsonify({"error": f"Audio processing failed: {str(e)}"}), 500

        # Transcribe user audio
        try:
            logger.info(f"[{request_id}] 🔄 Transcribing user audio")
            # Remove language parameter if causing warnings
            inputs = asr_processor(
                user_waveform,
                sampling_rate=sample_rate,
                return_tensors="pt"
            )
            inputs = {k: v.to(asr_model.device) for k, v in inputs.items()}

            with torch.no_grad():
                logits = asr_model(**inputs).logits
            ids = torch.argmax(logits, dim=-1)[0]
            user_transcription = asr_processor.decode(ids)

            logger.info(f"[{request_id}] ✅ User transcription: '{user_transcription}'")
        except Exception as e:
            logger.error(f"[{request_id}] ❌ ASR inference failed: {str(e)}")
            return jsonify({"error": f"ASR inference failed: {str(e)}"}), 500

        # Process reference files in batches
        batch_size = 2  # Process 2 files at a time - adjust based on your hardware
        results = []
        best_score = 0
        best_reference = None
        best_transcription = None

        # Use this if you want to limit the number of files to process
        max_files_to_check = min(5, len(reference_files))  # Check at most 5 files
        reference_files = reference_files[:max_files_to_check]

        logger.info(f"[{request_id}] 🔄 Processing {len(reference_files)} reference files in batches of {batch_size}")

        # Function to process a single reference file
        def process_reference_file(ref_file):
            ref_filename = os.path.basename(ref_file)
            try:
                # Load and resample reference audio
                ref_waveform, ref_sr = torchaudio.load(ref_file)
                if ref_sr != sample_rate:
                    ref_waveform = torchaudio.transforms.Resample(ref_sr, sample_rate)(ref_waveform)
                ref_waveform = ref_waveform.squeeze().numpy()
        
                # Transcribe reference audio - use the local asr_model and asr_processor
                # Remove language parameter if causing warnings
                inputs = asr_processor(
                    ref_waveform,
                    sampling_rate=sample_rate,
                    return_tensors="pt"
                )
                inputs = {k: v.to(asr_model.device) for k, v in inputs.items()}

                with torch.no_grad():
                    logits = asr_model(**inputs).logits
                ids = torch.argmax(logits, dim=-1)[0]
                ref_transcription = asr_processor.decode(ids)

                # Calculate similarity
                similarity = calculate_similarity(user_transcription, ref_transcription)

                logger.info(
                    f"[{request_id}] 📊 Similarity with {ref_filename}: {similarity:.2f}%, transcription: '{ref_transcription}'")

                return {
                    "reference_file": ref_filename,
                    "reference_text": ref_transcription,
                    "similarity_score": similarity
                }
            except Exception as e:
                logger.error(f"[{request_id}] ❌ Error processing {ref_filename}: {str(e)}")
                return {
                    "reference_file": ref_filename,
                    "reference_text": "Error",
                    "similarity_score": 0,
                    "error": str(e)
                }

        # Process files in batches using ThreadPoolExecutor
        with ThreadPoolExecutor(max_workers=batch_size) as executor:
            batch_results = list(executor.map(process_reference_file, reference_files))
            results.extend(batch_results)

            # Find the best result
            for result in batch_results:
                if result["similarity_score"] > best_score:
                    best_score = result["similarity_score"]
                    best_reference = result["reference_file"]
                    best_transcription = result["reference_text"]

                    # Exit early if we found a very good match (optional)
                    if best_score > 80.0:
                        logger.info(f"[{request_id}] 🏁 Found excellent match: {best_score:.2f}%")
                        break

        # Clean up temp files
        try:
            if temp_dir and os.path.exists(temp_dir):
                shutil.rmtree(temp_dir)
                logger.debug(f"[{request_id}] 🧹 Cleaned up temporary directory")
        except Exception as e:
            logger.warning(f"[{request_id}] ⚠️ Failed to clean up temp files: {str(e)}")

        # Determine feedback based on score
        is_correct = best_score >= 70.0

        if best_score >= 90.0:
            feedback = "Perfect pronunciation! Excellent job!"
        elif best_score >= 80.0:
            feedback = "Great pronunciation! Your accent is very good."
        elif best_score >= 70.0:
            feedback = "Good pronunciation. Keep practicing!"
        elif best_score >= 50.0:
            feedback = "Fair attempt. Try focusing on the syllables that differ from the sample."
        else:
            feedback = "Try again. Listen carefully to the sample pronunciation."

        logger.info(f"[{request_id}] 📊 Final evaluation results: score={best_score:.2f}%, is_correct={is_correct}")
        logger.info(f"[{request_id}] 📝 Feedback: '{feedback}'")
        logger.info(f"[{request_id}] ✅ Evaluation complete")

        # Sort results by score descending
        results.sort(key=lambda x: x["similarity_score"], reverse=True)

        return jsonify({
            "is_correct": is_correct,
            "score": best_score,
            "feedback": feedback,
            "user_transcription": user_transcription,
            "best_reference_transcription": best_transcription,
            "reference_locator": reference_locator,
            "details": results
        })

    except Exception as e:
        logger.error(f"[{request_id}] ❌ Unhandled exception in evaluation endpoint: {str(e)}")
        logger.debug(f"[{request_id}] Stack trace: {traceback.format_exc()}")

        # Clean up on error
        try:
            if temp_dir and os.path.exists(temp_dir):
                shutil.rmtree(temp_dir)
        except:
            pass

        return jsonify({"error": f"Internal server error: {str(e)}"}), 500