# 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")

# Initialize cache at module level instead
EVALUATION_CACHE = {}

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 without dummy files"""
    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", "adwa", "anam", "apat", "apulu", "atlu", "dalan", "libu", "lima",  
        "metung", "pitu", "siyam", "walu", "masala", "madalumdum", "maragul", "marimla", "malagu", "marok", "mababa", "malapit", "matuling", "maputi", 
        "arung", "asbuk", "balugbug", "bitis", "buntuk", "butit", "gamat", "kuku", "salu", "tud", 
        "pisan", "dara", "achi", "apu", "ima", "tatang", "pengari", "koya", "kapatad", "wali", 
        "pasbul", "awang", "dagis", "bale", "ulas", "sambra", "sulu", "pitudturan", "luklukan", "ulnan"
    ]
    
    created_dirs = 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
    
    return created_dirs, 0  # Return 0 created files since we're not creating dummy files anymore

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 WITHOUT dummy files
        dirs_created, _ = setup_reference_patterns(working_dir)
        logger.info(f"📊 Created {dirs_created} directories")
        
        # Try to copy reference files from other found directories to working directory if needed
        if 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
                        # But skip any dummy reference files
                        for item in os.listdir(source_dir):
                            src_path = os.path.join(source_dir, item)
                            if os.path.isdir(src_path):
                                wav_files = glob.glob(os.path.join(src_path, "*.wav"))
                                # Filter out dummy references
                                wav_files = [f for f in wav_files if "dummy_reference" not in f]
                                
                                if wav_files:  # Only proceed if there are valid files
                                    dst_path = os.path.join(working_dir, item)
                                    os.makedirs(dst_path, exist_ok=True)
                                    
                                    # Copy each valid WAV file individually
                                    for wav_file in wav_files:
                                        wav_name = os.path.basename(wav_file)
                                        if "dummy_reference" not in wav_name:  # Extra check
                                            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, excluding dummy files - MODIFIED HERE
        pattern_dirs = [d for d in os.listdir(working_dir)
                       if os.path.isdir(os.path.join(working_dir, d))]
        
        # Count total files without logging each directory
        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"))
            # Count only non-dummy files
            valid_files = [f for f in wav_files if "dummy_reference" not in f]
            total_wav_files += len(valid_files)
            # Remove the individual directory logging
            # logger.info(f"  - {pattern}: {len(valid_files)} valid WAV files")
        
        logger.info(f"📊 Total pattern directories: {len(pattern_dirs)}, Total reference WAV files: {total_wav_files}")
        
        # Check for and remove any dummy files
        for pattern in pattern_dirs:
            pattern_path = os.path.join(working_dir, pattern)
            dummy_files = glob.glob(os.path.join(pattern_path, "dummy_reference.wav"))
            for dummy in dummy_files:
                try:
                    os.remove(dummy)
                    logger.info(f"🗑️ Removed dummy file: {dummy}")
                except Exception as e:
                    logger.warning(f"⚠️ Failed to remove dummy file {dummy}: {str(e)}")
        
        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 but without dummy files
        fallback_dir = os.path.join('/tmp', 'reference_audios')
        try:
            os.makedirs(fallback_dir, exist_ok=True)
            setup_reference_patterns(fallback_dir)  # This now doesn't create dummy files
            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", "adwa", "anam", "apat", "apulu", "atlu", "dalan", "libu", "lima",  
            "metung", "pitu", "siyam", "walu", "masala", "madalumdum", "maragul", "marimla", "malagu", "marok", "mababa", "malapit", "matuling", "maputi", 
            "arung", "asbuk", "balugbug", "bitis", "buntuk", "butit", "gamat", "kuku", "salu", "tud", 
            "pisan", "dara", "achi", "apu", "ima", "tatang", "pengari", "koya", "kapatad", "wali", 
            "pasbul", "awang", "dagis", "bale", "ulas", "sambra", "sulu", "pitudturan", "luklukan", "ulnan"
        ]

        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 with speed optimizations"""
    request_id = f"req-{id(request)}"
    logger.info(f"[{request_id}] 🆕 Starting 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:
        # OPTIMIZATION 1: Check cache first for identical audio
        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

        # OPTIMIZATION 2: Simple caching based on audio content hash + reference_locator
        audio_content = audio_file.read()
        audio_file.seek(0)  # Reset file pointer after reading
        
        import hashlib
        audio_hash = hashlib.md5(audio_content).hexdigest()
        cache_key = f"{audio_hash}_{reference_locator}_{language}"
        
        # Check in-memory cache using the module-level cache
        if cache_key in EVALUATION_CACHE:
            logger.info(f"[{request_id}] ✅ Using cached evaluation result")
            return EVALUATION_CACHE[cache_key]

        # 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"))
        # Filter out any dummy reference files
        reference_files = [f for f in reference_files if "dummy_reference" not in f]
        logger.info(f"[{request_id}] 📁 Found {len(reference_files)} valid reference files")

        # If no reference files exist, return a more detailed error message
        if not reference_files:
            logger.warning(f"[{request_id}] ⚠️ No valid reference audio files found in {reference_dir_path}")
            return jsonify({
                "error": f"No reference audio found for {reference_locator}",
                "message": "Please upload a reference audio file before evaluation.",
                "status": "MISSING_REFERENCE"
            }), 404

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

        # 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_content)  # Use the content we already 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

        # OPTIMIZATION 3: Use a smaller sample of reference files
        import multiprocessing
        import random
        
        # OPTIMIZATION 4: Limit to just a few files for initial comparison
        # If we have many reference files, randomly sample some for quick evaluation
        if len(reference_files) > 3:
            # Randomly select 3 files for faster comparison
            reference_files_sample = random.sample(reference_files, 3)
        else:
            reference_files_sample = reference_files
            
        # Determine optimal number of workers based on CPU count (but keep it small)
        max_workers = min(multiprocessing.cpu_count(), len(reference_files_sample), 3)
        initial_results = []

        logger.info(f"[{request_id}] 🔄 Quick scan: processing {len(reference_files_sample)} reference files")

        # 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
                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 the sample files in parallel
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            initial_results = list(executor.map(process_reference_file, reference_files_sample))
        
        # Find the best result from the initial sample
        best_score = 0
        best_reference = None
        best_transcription = None
        
        for result in initial_results:
            if result["similarity_score"] > best_score:
                best_score = result["similarity_score"]
                best_reference = result["reference_file"]
                best_transcription = result["reference_text"]
        
        # OPTIMIZATION 5: If we already found a very good match, don't process more files
        all_results = initial_results.copy()
        remaining_files = [f for f in reference_files if f not in reference_files_sample]
        
        # Only process more files if our best score isn't already very good
        if best_score < 80.0 and remaining_files:
            logger.info(f"[{request_id}] 🔄 Score {best_score:.2f}% not high enough, checking {len(remaining_files)} more references")
            
            # Process remaining files
            with ThreadPoolExecutor(max_workers=max_workers) as executor:
                additional_results = list(executor.map(process_reference_file, remaining_files[:5]))  # Process max 5 more
                all_results.extend(additional_results)
                
                # Update best result if we found a better one
                for result in additional_results:
                    if result["similarity_score"] > best_score:
                        best_score = result["similarity_score"]
                        best_reference = result["reference_file"]
                        best_transcription = result["reference_text"]

        # 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
        all_results.sort(key=lambda x: x["similarity_score"], reverse=True)
        
        # Create response
        response = jsonify({
            "is_correct": is_correct,
            "score": best_score,
            "feedback": feedback,
            "user_transcription": user_transcription,
            "best_reference_transcription": best_transcription,
            "reference_locator": reference_locator,
            "details": all_results,
            "total_references_compared": len(all_results),
            "total_available_references": len(reference_files),
            "quick_evaluation": True
        })
        
        # OPTIMIZATION 6: Cache the result for future requests using module-level cache
        MAX_CACHE_SIZE = 50
        EVALUATION_CACHE[cache_key] = response
        if len(EVALUATION_CACHE) > MAX_CACHE_SIZE:
            # Remove oldest entry (simplified approach)
            EVALUATION_CACHE.pop(next(iter(EVALUATION_CACHE)))
        
        return response

    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