import os
import sys
import time
from datetime import datetime
from google import genai
from google.genai import types
from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type
import pysubs2
from deep_translator import GoogleTranslator, MyMemoryTranslator
import ollama
from tqdm import tqdm
from utils import LANGUAGE_MAP

# Define a retry decorator
# ... (retry_policy remains)

def translate_via_ollama(source_srt_content, target_language="English", model="dolphin-llama3"):
    """
    Translates SRT content using a local Ollama model (Line-by-Line for progress).
    Includes retries and debug logging.
    """
    debug_log_path = "ollama_debug.log"
    
    try:
        subs = pysubs2.SSAFile.from_string(source_srt_content)
        
        # Using tqdm for progress bar
        # dynamic_ncols=True helps it resize properly.
        for line in tqdm(subs, desc="  Ollama Progress", unit="line", dynamic_ncols=True, leave=False):
            text = line.text.strip()
            
            # Skip empty, numeric-only, or extremely short non-word text
            if not text or text.isdigit() or len(text) < 2:
                continue

            prompt = (
                f"Translate this subtitle text to {target_language}. Output ONLY the translation.\n"
                f"Text: {text}"
            )
            
            # Retry loop for stability
            max_retries = 3
            for attempt in range(max_retries):
                try:
                    response = ollama.chat(model=model, messages=[{'role': 'user', 'content': prompt}])
                    
                    # Check for "model not found" or other soft errors in response if API wraps them
                    # Usually ollama library raises ResponseError for 404
                    
                    translated_text = response['message']['content'].strip()
                    if translated_text:
                        line.text = translated_text
                    break # Success, exit retry loop
                    
                except Exception as e:
                    # Log error details
                    with open(debug_log_path, "a") as log:
                        log.write(f"[{datetime.now()}] Error on line '{text}': {str(e)}\n")
                    
                    if attempt < max_retries - 1:
                        time.sleep(2) # Wait before retry
                    else:
                        # If all retries fail, keep original text or empty?
                        # Keeping original might be safer than silence, or just skip.
                        # For now, we skip updating 'line.text' so it stays as source language (better than corruption)
                        pass

        return subs.to_string(format_="srt")

    except Exception as e:
        print(f"  [Local LLM] Critical Error: {e}")
        return None

def translate_fallback_mymemory(source_srt_content, target_language="en"):
    """
    Fallback translation using MyMemory (via deep-translator).
    Limit: 1000 words/day roughly for anonymous usage. Good last resort.
    """
    try:
        subs = pysubs2.SSAFile.from_string(source_srt_content)
        # MyMemory uses ISO 639-1 usually
        translator = MyMemoryTranslator(source='auto', target=target_language)
        
        for line in tqdm(subs, desc="  MyMemory Progress", unit="line", leave=False):
            text = line.text.strip()
            
            # Skip empty, numeric-only, or extremely short non-word text
            if not text or text.isdigit() or len(text) < 2:
                continue

            if text:
                if len(text) > 500: # MyMemory has stricter limits often
                     continue
                try:
                    original_text = text.replace(r"\N", " ")
                    translated_text = translator.translate(original_text)
                    if translated_text:
                        line.text = translated_text
                except Exception:
                    pass
        
        return subs.to_string(format_="srt")
    except Exception as e:
        print(f"  [MyMemory Fallback] Critical Error: {e}")
        return None

def translate_fallback_free(source_srt_content, target_language="en"):
    """
    Fallback translation using deep-translator (free Google Translate).
    
    Args:
        source_srt_content (str): Content of the source SRT file.
        target_language (str): Target language code (e.g. 'en', 'fr').
    
    Returns:
        str: Translated SRT content, or None if failed.
    """
    try:
        # Load from string
        subs = pysubs2.SSAFile.from_string(source_srt_content)
        translator = GoogleTranslator(source='auto', target=target_language)
        
        # Simple line-by-line translation
        for line in tqdm(subs, desc="  DeepTranslate Progress", unit="line", leave=False):
            text = line.text.strip()
            
            # Skip empty, numeric-only, or extremely short non-word text
            if not text or text.isdigit() or len(text) < 2:
                continue

            if text:
                # Sanity check: Skip lines that are too long
                if len(text) > 4000:
                    continue

                try:
                    # pysubs2 text can contain \N for newlines. 
                    original_text = text.replace(r"\N", " ")
                    translated_text = translator.translate(original_text)
                    if translated_text:
                        line.text = translated_text
                except Exception:
                    pass
        
        # Return as string
        return subs.to_string(format_="srt")
    except Exception as e:
        print(f"  [Free Fallback] Critical Error: {e}")
        return None

# Define a retry decorator
# Waits 2^x * 1 seconds between retries (1s, 2s, 4s...)
# Stop after 15 attempts
# before_sleep logic can print a simple message
def log_retry_attempt(retry_state):
    if retry_state.attempt_number > 1:
        print(f"  [Gemini] Rate limit hit. Retrying in {retry_state.next_action.sleep}s...", end='\r')

retry_policy = retry(
    stop=stop_after_attempt(15),
    wait=wait_exponential(multiplier=1, min=2, max=60),
    retry=retry_if_exception_type(Exception),
    reraise=True,
    before_sleep=log_retry_attempt
)

@retry_policy
def _generate_with_retry(client, model_name, prompt):
    """Internal function to wrap the API call with retry logic."""
    return client.models.generate_content(
        model=model_name,
        contents=prompt
    )

def get_best_available_model(client):
    """
    Queries the API to find the best available model for text generation.
    Priority: gemini-2.0-flash > gemini-1.5-flash > gemini-1.5-pro
    """
    try:
        # Priority list (New v2 naming conventions if applicable, but standard models persist)
        priorities = [
            "gemini-2.0-flash", # Latest
            "gemini-1.5-flash",
            "gemini-1.5-pro"
        ]
        
        # In new SDK, client.models.list() returns iterators of Model objects
        # We can just try to use the priority one directly, or list them.
        # Listing can be slow. Let's just default to a known good priority list.
        # If we really want to check:
        # available = [m.name for m in client.models.list()]
        
        # For efficiency/speed, we will trust our priority list. 
        # The API will error if model doesn't exist, which the try/catch block handling generation will catch?
        # No, better to pick one that exists.
        
        # Let's return the latest standard one.
        return "gemini-2.0-flash" # Assuming 2.0 is available or falling back
            
    except Exception as e:
        print(f"Warning: Model selection issue ({e}). Defaulting to 'gemini-1.5-flash'.")
    
    return "gemini-1.5-flash"

def translate_srt(srt_content, target_language="English", api_key=None):
    """
    Translates SRT subtitle content using the Google GenAI SDK (v2).
    """
    if not srt_content:
        return ""

    key = api_key or os.getenv("GEMINI_API_KEY")
    if not key:
        print("Error: GEMINI_API_KEY not found. Please set the environment variable or pass the key.")
        sys.exit(1)

    # Initialize Client (v2 style)
    try:
        client = genai.Client(api_key=key)
    except Exception as e:
        print(f"Error initializing GenAI Client: {e}")
        return None

    # Automatically select the best model
    # Note: v2 SDK might use 'gemini-1.5-flash' directly without 'models/' prefix usually
    model_name = "gemini-2.0-flash" 

    prompt = (
        "You are a professional subtitle translator. Your task is to translate the following SRT subtitle file "
        f"into {target_language}.\n\n"
        "RULES:\n"
        "1. PRESERVE the SRT format exactly. Do not modify timestamps (e.g., 00:00:01,000 --> 00:00:04,000) or sequence numbers.\n"
        "2. Only translate the dialogue text.\n"
        "3. Maintain the original tone and context.\n"
        "4. Output ONLY the translated SRT content, no markdown code blocks or explanations.\n\n"
        "SRT Content:\n"
        f"{srt_content}"
    )

    try:
        # Call the retried internal function
        response = _generate_with_retry(client, model_name, prompt)
        
        # Cleanup: sometimes models wrap output in ```srt ... ``` or ``` ... ```
        cleaned_text = response.text.strip()
        if cleaned_text.startswith("```"):
            lines = cleaned_text.split('\n')
            if len(lines) >= 2:
                cleaned_text = '\n'.join(lines[1:-1])
        
        return cleaned_text
    except Exception as e:
        print(f"Error during translation after retries: {e}")
        # Fallback to older model if 2.0 fails?
        if "404" in str(e) and "gemini-2.0" in model_name:
             print("  -> gemini-2.0-flash not found, falling back to gemini-1.5-flash")
             try:
                 response = _generate_with_retry(client, "gemini-1.5-flash", prompt)
                 cleaned_text = response.text.strip()
                 if cleaned_text.startswith("```"):
                    lines = cleaned_text.split('\n')
                    if len(lines) >= 2:
                        cleaned_text = '\n'.join(lines[1:-1])
                 return cleaned_text
             except Exception as inner_e:
                 print(f"Fallback failed: {inner_e}")

        return None

def translate_with_auto_fallback(srt_content, target_language="English", prefer_deep=False, prefer_local=False, available_services=None):
    """
    Attempts to translate SRT content using Gemini, DeepTranslate, and Local LLM with fallback logic.
    
    Args:
        srt_content (str): The source SRT content.
        target_language (str): Target language name (e.g., "English", "French").
        prefer_deep (bool): If True, try DeepTranslate first (among cloud services).
        prefer_local (bool): If True, try Local LLM (Ollama) first.
        available_services (dict, optional): Result of check_service_availability().
        
    Returns:
        tuple: (translated_content, method_name) or (None, None) if all failed.
    """
    
    # Use central language mapping
    target_code = LANGUAGE_MAP.get(target_language, "en")

    # Determine which services to even try
    def is_ok(name):
        if available_services is None: return True
        return available_services.get(name, True)

    def try_gemini():
        if not is_ok("Gemini"): return None, None
        res = translate_srt(srt_content, target_language=target_language)
        if res: return res, "Gemini"
        return None, None

    def try_deep():
        if not is_ok("DeepTranslate"): return None, None
        res = translate_fallback_free(srt_content, target_language=target_code)
        if res: return res, "DeepTranslate"
        return None, None

    def try_ollama():
        if not is_ok("Ollama"): return None, None
        res = translate_via_ollama(srt_content, target_language=target_language)
        if res: return res, "Local LLM (Ollama)"
        return None, None

    def try_mymemory():
        res = translate_fallback_mymemory(srt_content, target_language=target_code)
        if res: return res, "MyMemory"
        return None, None

    # Logic flow
    attempts = []
    
    if prefer_local:
        attempts.append(try_ollama)
        if prefer_deep:
            attempts.extend([try_deep, try_gemini])
        else:
            attempts.extend([try_gemini, try_deep])
    else:
        if prefer_deep:
            attempts.extend([try_deep, try_gemini])
        else:
            attempts.extend([try_gemini, try_deep])
        attempts.append(try_ollama)
    
    # Final last resort
    attempts.append(try_mymemory)

    # Execute attempts
    for i, method_func in enumerate(attempts):
        # if i > 0:
        #    print(f"  Attempt {i} failed. Trying next fallback...")
        
        content, method = method_func()
        if content:
            return content, method

    return None, None