import asyncio
from datetime import datetime
from typing import Any
import time
import random

from celery import Task
from motor.motor_asyncio import AsyncIOMotorClient

from ..core.config import settings
from ..core.logging import get_logger
from ..models.job import JobStatus
from ..services.gcs import gcs_service, upload_vtt_to_gcs
from ..services.gemini import gemini_service
from ..services.translate import translate_service
from ..services.tts import tts_service
from ..services.websocket import connection_manager
from . import celery_app

logger = get_logger(__name__)


def broadcast_status_update(job_id: str, status: str, job_title: str = None, message: str = None, progress: int = None):
    """
    Helper function to broadcast job status updates via WebSocket
    Uses sync Redis client for Celery worker context
    """
    logger.info(f"🔊 ATTEMPTING TO BROADCAST: job_id={job_id}, status={status}, job_title={job_title}, message={message}")
    try:
        import redis as sync_redis
        from ..core.config import settings
        from ..services.websocket import JobStatusUpdate
        from datetime import datetime
        
        logger.info(f"🔊 About to create JobStatusUpdate for job {job_id}")
        
        # Create status update
        update = JobStatusUpdate(
            job_id=job_id,
            status=status,
            updated_at=datetime.utcnow(),
            job_title=job_title,
            message=message,
            progress=progress
        )
        
        logger.info(f"🔊 Created update object, now connecting to Redis: {settings.redis_url}")
        
        # Create synchronous Redis client
        redis_client = sync_redis.Redis.from_url(
            settings.redis_url,
            encoding="utf-8",
            decode_responses=True
        )
        
        logger.info(f"🔊 Redis client created, now publishing to channels")
        
        # Publish to channels
        result1 = redis_client.publish("job_status_updates", update.model_dump_json())
        result2 = redis_client.publish(f"job_status_updates:{job_id}", update.model_dump_json())
        
        logger.info(f"🔊 Published to channels - general: {result1} subscribers, job-specific: {result2} subscribers")
        
        # Close connection
        redis_client.close()
        
        logger.info(f"🔊 ✅ Successfully broadcasted status update for job {job_id}: {status}")
        
    except Exception as e:
        logger.error(f"🔊 ❌ Failed to broadcast status update for job {job_id}: {e}")
        import traceback
        logger.error(f"🔊 ❌ Full traceback: {traceback.format_exc()}")


async def retry_with_backoff(func, max_retries=3, base_delay=1):
    """Retry a function with exponential backoff"""
    last_exception = None
    
    for attempt in range(max_retries):
        try:
            return await func()
        except Exception as e:
            last_exception = e
            if attempt == max_retries - 1:
                break
                
            # Exponential backoff with jitter
            delay = base_delay * (2 ** attempt) + random.uniform(0, 1)
            logger.warning(f"Attempt {attempt + 1} failed, retrying in {delay:.2f}s: {e}")
            await asyncio.sleep(delay)
    
    raise last_exception


@celery_app.task(bind=True)
def translate_and_synthesize_task(self, job_id: str):
    """
    Pipeline 2: Translation & MP3 Generation
    Triggered when job status changes to 'approved_english'
    """
    logger.info(f"🚀 CELERY TASK STARTED: translate_and_synthesize_task for job {job_id}")
    
    try:
        logger.info(f"📝 About to call asyncio.run for job {job_id}")
        result = asyncio.run(_async_translate_and_synthesize(job_id))
        logger.info(f"✅ CELERY TASK COMPLETED successfully for job {job_id}")
        return result
    except Exception as e:
        logger.error(f"❌ CELERY TASK FAILED for job {job_id}: {str(e)}")
        logger.error(f"❌ Exception type: {type(e).__name__}")
        logger.error(f"❌ Exception args: {e.args}")
        import traceback
        logger.error(f"❌ Full traceback: {traceback.format_exc()}")
        raise


async def _async_translate_and_synthesize(job_id: str):
    """Async implementation of translation and synthesis"""
    logger.info(f"🔄 ASYNC FUNCTION STARTED: _async_translate_and_synthesize for job {job_id}")
    
    # Connect to MongoDB
    logger.info(f"📡 Connecting to MongoDB for job {job_id}")
    client = AsyncIOMotorClient(settings.mongodb_uri)
    db = client[settings.mongodb_db]
    logger.info(f"📡 MongoDB connection established for job {job_id}")

    try:
        # Get job details
        logger.info(f"🔍 Looking up job document for job {job_id}")
        job_doc = await db.jobs.find_one({"_id": job_id})
        if not job_doc:
            logger.error(f"❌ Job {job_id} not found in database!")
            raise ValueError(f"Job {job_id} not found")

        job_title = job_doc.get("title", "Untitled Job")
        logger.info(f"✅ Found job document for {job_id} ({job_title}), status: {job_doc.get('status', 'UNKNOWN')}")

        # Check for any approved status (English or non-English source)
        if not JobStatus.is_approved(job_doc["status"]):
            logger.warning(f"⚠️ Job {job_id} not in approved status (current: {job_doc['status']}), skipping")
            return

        # Get source language from job
        source_language = job_doc["source"].get("language", "en")
        logger.info(f"✅ Job {job_id} is in correct status, proceeding with translation (source: {source_language})")

        # Update status to translating
        await db.jobs.update_one(
            {"_id": job_id},
            {
                "$set": {
                    "status": JobStatus.TRANSLATING.value,
                    "updated_at": datetime.utcnow()
                },
                "$push": {
                    "review.history": {
                        "at": datetime.utcnow(),
                        "status": JobStatus.TRANSLATING.value,
                        "by": "system"
                    }
                }
            }
        )

        # Get source language VTT content
        source_outputs = job_doc["outputs"].get(source_language)
        if not source_outputs:
            raise ValueError(f"No outputs found for source language {source_language}")

        # Download source language VTT files
        captions_blob_path = source_outputs["captions_vtt_gcs"].replace(f"gs://{settings.gcs_bucket}/", "")
        ad_blob_path = source_outputs["ad_vtt_gcs"].replace(f"gs://{settings.gcs_bucket}/", "")

        captions_blob = gcs_service.bucket.blob(captions_blob_path)
        ad_blob = gcs_service.bucket.blob(ad_blob_path)

        source_captions_vtt = captions_blob.download_as_text()
        source_ad_vtt = ad_blob.download_as_text()

        # Process each requested language
        requested_languages = job_doc["requested_outputs"]["languages"]
        transcreation_languages = job_doc["requested_outputs"]["transcreation"]

        updated_outputs = job_doc.get("outputs", {})

        for language in requested_languages:
            if language == source_language:
                continue  # Skip source language as it's already processed

            logger.info(f"Processing language: {language} (from source: {source_language})")

            try:
                if language in transcreation_languages:
                    # Use transcreation for cultural adaptation with retry
                    async def transcreate():
                        return await gemini_service.transcreate_content(
                            source_captions_vtt,
                            source_ad_vtt,
                            language,
                            brief="Standard accessibility content"
                        )

                    result = await retry_with_backoff(transcreate, max_retries=3)
                    translated_captions = result["captions_vtt"]
                    translated_ad = result["audio_description_vtt"]
                    origin = "transcreate"

                else:
                    # Use standard translation with retry
                    async def translate_captions():
                        return await translate_service.translate_vtt(
                            source_captions_vtt, language, source_language=source_language
                        )

                    async def translate_ad():
                        return await translate_service.translate_vtt(
                            source_ad_vtt, language, source_language=source_language
                        )

                    translated_captions = await retry_with_backoff(translate_captions, max_retries=3)
                    translated_ad = await retry_with_backoff(translate_ad, max_retries=3)
                    origin = "translate"

                # Upload translated VTT files
                captions_gcs_uri = await upload_vtt_to_gcs(
                    translated_captions,
                    f"{job_id}/{language}/captions.vtt"
                )

                ad_gcs_uri = await upload_vtt_to_gcs(
                    translated_ad,
                    f"{job_id}/{language}/ad.vtt"
                )

                # Store language outputs
                updated_outputs[language] = {
                    "captions_vtt_gcs": captions_gcs_uri,
                    "ad_vtt_gcs": ad_gcs_uri,
                    "origin": origin
                }

                logger.info(f"Successfully processed VTT files for language: {language}")

            except Exception as e:
                logger.error(f"Failed to process language {language}: {e}")
                updated_outputs[language] = {
                    "origin": "translate" if language not in transcreation_languages else "transcreate",
                    "qa_notes": f"Translation failed: {str(e)}"
                }

        # Update status to TTS generating
        await db.jobs.update_one(
            {"_id": job_id},
            {
                "$set": {
                    "status": JobStatus.TTS_GENERATING.value,
                    "outputs": updated_outputs,
                    "updated_at": datetime.utcnow()
                },
                "$push": {
                    "review.history": {
                        "at": datetime.utcnow(),
                        "status": JobStatus.TTS_GENERATING.value,
                        "by": "system"
                    }
                }
            }
        )

        # Generate TTS for languages that need MP3
        if job_doc["requested_outputs"]["audio_description_mp3"]:
            # Get TTS preferences from job
            tts_preferences = job_doc["requested_outputs"].get("tts_preferences", {})
            await _generate_tts_for_languages(job_id, updated_outputs, db, source_language, tts_preferences)

        # Update final status
        await db.jobs.update_one(
            {"_id": job_id},
            {
                "$set": {
                    "status": JobStatus.PENDING_FINAL_REVIEW.value,
                    "updated_at": datetime.utcnow()
                },
                "$push": {
                    "review.history": {
                        "at": datetime.utcnow(),
                        "status": JobStatus.PENDING_FINAL_REVIEW.value,
                        "by": "system"
                    }
                }
            }
        )
        
        # Broadcast status update
        broadcast_status_update(
            job_id,
            JobStatus.PENDING_FINAL_REVIEW.value,
            job_title=job_title,
            message=f"{job_title} has finished translation and audio generation - ready for Final Review"
        )

        logger.info(f"Successfully completed translation and synthesis for job {job_id}")

    except Exception as e:
        logger.error(f"Translation and synthesis failed for job {job_id}: {e}")

        # Update job with error
        await db.jobs.update_one(
            {"_id": job_id},
            {
                "$set": {
                    "error": {
                        "type": "translation_failure",
                        "message": str(e),
                        "timestamp": datetime.utcnow().isoformat()
                    },
                    "updated_at": datetime.utcnow()
                }
            }
        )

        raise

    finally:
        client.close()


async def _generate_tts_for_languages(
    job_id: str,
    outputs: dict[str, Any],
    db,
    source_language: str = "en",
    tts_preferences: dict = None
):
    """Generate TTS audio for each language's audio description"""
    if tts_preferences is None:
        tts_preferences = {}

    # Always generate source language MP3 first
    if source_language in outputs and "ad_vtt_gcs" in outputs[source_language]:
        await _generate_language_tts(job_id, source_language, outputs[source_language], db, tts_preferences)

    # Generate for other languages
    for language, lang_output in outputs.items():
        if language != source_language and "ad_vtt_gcs" in lang_output:
            await _generate_language_tts(job_id, language, lang_output, db, tts_preferences)


async def _generate_language_tts(job_id: str, language: str, lang_output: dict, db, tts_preferences: dict = None):
    """Generate TTS for a specific language"""
    if tts_preferences is None:
        tts_preferences = {}

    try:
        # Download AD VTT content
        ad_blob_path = lang_output["ad_vtt_gcs"].replace(f"gs://{settings.gcs_bucket}/", "")
        ad_blob = gcs_service.bucket.blob(ad_blob_path)
        ad_vtt_content = ad_blob.download_as_text()

        # Get voice for this language from preferences
        voices_per_language = tts_preferences.get("voices_per_language", {})
        voice_name = voices_per_language.get(language, tts_preferences.get("default_voice"))
        provider = tts_preferences.get("provider", "gemini")

        # Get new TTS settings with defaults
        model = tts_preferences.get("model", "flash")
        speed = tts_preferences.get("speed", 1.0)
        style_preset = tts_preferences.get("style_preset", "neutral")
        custom_style_prompt = tts_preferences.get("custom_style_prompt")

        # Resolve style prompt from preset or custom
        if style_preset == "custom" and custom_style_prompt:
            style_prompt = custom_style_prompt
        else:
            style_prompt = settings.gemini_tts_style_prompts.get(style_preset, "")

        # Generate MP3 with retry
        language_code = f"{language}-US" if language == "en" else f"{language}-{language.upper()}"

        logger.info(
            f"Generating TTS for {language} with voice={voice_name}, provider={provider}, "
            f"model={model}, speed={speed}x, style={style_preset}"
        )

        async def synthesize():
            return await tts_service.synthesize_audio_description(
                ad_vtt_content,
                language_code,
                voice_name=voice_name,
                provider=provider,
                model=model,
                speed=speed,
                style_prompt=style_prompt
            )

        mp3_data = await retry_with_backoff(synthesize, max_retries=3)

        # Upload MP3 to GCS
        mp3_blob_path = f"{job_id}/{language}/ad.mp3"
        mp3_blob = gcs_service.bucket.blob(mp3_blob_path)
        mp3_blob.content_type = "audio/mpeg"
        mp3_blob.upload_from_string(mp3_data, content_type="audio/mpeg")

        mp3_gcs_uri = f"gs://{settings.gcs_bucket}/{mp3_blob_path}"

        # Update job outputs
        await db.jobs.update_one(
            {"_id": job_id},
            {
                "$set": {
                    f"outputs.{language}.ad_mp3_gcs": mp3_gcs_uri,
                    "updated_at": datetime.utcnow()
                }
            }
        )

        logger.info(f"Successfully generated TTS for {language}")

        # Broadcast asset ready notification immediately
        # This allows frontend to update in real-time as each MP3 becomes available
        job_doc = await db.jobs.find_one({"_id": job_id})
        broadcast_status_update(
            job_id,
            "asset_ready",
            job_title=job_doc.get("title"),
            message=f"Audio description MP3 ready for {language.upper()}"
        )

    except Exception as e:
        logger.error(f"TTS generation failed for {language}: {e}")

        # Update with error note
        await db.jobs.update_one(
            {"_id": job_id},
            {
                "$set": {
                    f"outputs.{language}.qa_notes": f"TTS generation failed: {str(e)}",
                    "updated_at": datetime.utcnow()
                }
            }
        )