import asyncio
import os
import subprocess
import tempfile
from datetime import datetime

import ffmpeg
from celery import Task
from celery.result import allow_join_result
from motor.motor_asyncio import AsyncIOMotorClient

from ..core.config import settings
from ..core.logging import get_logger
from ..lib.vtt import VTTEditor
from ..models.job import JobStatus
from ..services import caption_aligner, cost_tracker
from ..services.gcs import gcs_path, gcs_service, upload_vtt_to_gcs
from ..services.gemini import gemini_service
from ..services.whisper_service import WordTimestamp
from . import celery_app
from ._websocket_bridge import broadcast_status_update
from .whisper_transcribe import transcribe_video_audio_task

logger = get_logger(__name__)


class AsyncTask(Task):
    """Base task class that supports async execution"""
    def __call__(self, *args, **kwargs):
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        try:
            return loop.run_until_complete(self.run_async(*args, **kwargs))
        finally:
            loop.close()

    async def run_async(self, *args, **kwargs):
        raise NotImplementedError


class IngestAndAITask(AsyncTask):
    async def run_async(self, job_id: str):
        return await ingest_and_ai_task_impl(job_id)


@celery_app.task(bind=True, base=IngestAndAITask)
def ingest_and_ai_task(self, job_id: str):
    """
    Pipeline 1: Ingestion & AI Processing
    Task wrapper that delegates to async implementation
    """
    # This method is called by AsyncTask.__call__
    pass


async def ingest_and_ai_task_impl(job_id: str):
    """
    Pipeline 1: Ingestion & AI Processing
    1. Update status to 'ingesting'
    2. Probe video for metadata (duration, codec)
    3. Process with Gemini 2.5 Pro
    4. Generate VTT files
    5. Update status to 'pending_qc'
    """
    logger.info(f"Starting ingestion and AI processing for job {job_id}")

    # Connect to MongoDB
    client = AsyncIOMotorClient(settings.mongodb_uri)
    db = client[settings.mongodb_db]

    try:
        # Get job document to retrieve title for notifications
        job_doc = await db.jobs.find_one({"_id": job_id})
        if not job_doc:
            logger.error(f"Job {job_id} not found in database")
            return

        job_title = job_doc.get("title", "Untitled Job")
        logger.info(f"Processing job: {job_title}")

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

        # Broadcast ingesting status update
        broadcast_status_update(
            job_id,
            JobStatus.INGESTING.value,
            job_title=job_title,
            message=f"{job_title} is being ingested"
        )

        # Get job details
        job_doc = await db.jobs.find_one({"_id": job_id})
        if not job_doc:
            raise ValueError(f"Job {job_id} not found")

        # Download video file temporarily for processing
        source_blob_path = job_doc["source"]["gcs_uri"].replace(f"gs://{settings.gcs_bucket}/", "")

        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as temp_file:
            temp_path = temp_file.name

            # Download from GCS
            blob = gcs_service.bucket.blob(source_blob_path)
            blob.download_to_filename(temp_path)

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

                # Broadcast AI processing status update
                broadcast_status_update(
                    job_id,
                    JobStatus.AI_PROCESSING.value,
                    job_title=job_title,
                    message=f"{job_title} is being processed by AI"
                )

                # Probe video for metadata
                duration = await _get_video_duration(temp_path)

                # Update source with duration
                await db.jobs.update_one(
                    {"_id": job_id},
                    {"$set": {"source.duration_s": duration}}
                )

                # Process with Gemini
                brand_context = job_doc.get("brand_context")
                sdh_requested = job_doc.get("requested_outputs", {}).get("sdh_vtt", False)
                source_has_ad = job_doc.get("source", {}).get("source_has_ad", False)
                _cost_ctx = {
                    "user_id": job_doc.get("client_id", "system"),
                    "job_id": job_id,
                    "project_id": job_doc.get("cost_tracker_project_id"),
                }
                await cost_tracker.aio_preflight(
                    model=gemini_service.model_name,
                    user_external_id=_cost_ctx["user_id"],
                    project_id=_cost_ctx["project_id"],
                )
                # Load glossary for source language — use title + brand context for term matching
                from ..services.glossary_service import get_glossary_block_for_job
                _source_lang = job_doc.get("source", {}).get("language", "en")
                _job_title = job_doc.get("title") or ""
                _source_for_glossary = " ".join(filter(None, [_job_title, brand_context]))
                _job_for_glossary = {**job_doc, "_glossary_source_text": _source_for_glossary}
                glossary_block = await get_glossary_block_for_job(_job_for_glossary, _source_lang, db)
                ai_result = await gemini_service.extract_accessibility(
                    temp_path,
                    brand_context=brand_context,
                    sdh_requested=sdh_requested,
                    source_has_ad=source_has_ad,
                    glossary_block=glossary_block,
                    _cost_ctx=_cost_ctx,
                )
                # Enforce: if source already has AD, discard any AI-generated AD
                if source_has_ad:
                    ai_result["audio_description_vtt"] = "WEBVTT\n"
                    logger.info(f"source_has_ad=True for job {job_id}: skipping AD generation")

                # Final safety check for required fields
                required_fields = ["captions_vtt", "audio_description_vtt"]
                missing_fields = [field for field in required_fields if field not in ai_result]

                if missing_fields:
                    logger.error(f"Missing required fields after AI processing: {missing_fields}")
                    # Create fallback content for missing fields
                    if "audio_description_vtt" in missing_fields:
                        ai_result["audio_description_vtt"] = "WEBVTT\n\n00:00:00.000 --> 00:00:05.000\nVideo content with visual elements."
                        logger.info("Created fallback audio_description_vtt")

                # Get detected language from Gemini response
                detected_language = ai_result.get("language", "en")
                language_hint = job_doc["source"].get("language_hint")
                initial_language = job_doc["source"].get("language", "en")

                # Log if there's a mismatch between hint/initial and detected language
                if language_hint and language_hint != detected_language:
                    logger.warning(
                        f"Language mismatch for job {job_id}: "
                        f"hint={language_hint}, detected={detected_language}"
                    )
                elif initial_language != "auto" and initial_language != detected_language:
                    logger.info(
                        f"Language detection for job {job_id}: "
                        f"initial={initial_language}, detected={detected_language}"
                    )

                # Use detected language for output storage
                source_language = detected_language
                logger.info(f"Using detected language '{source_language}' for job {job_id}")

                # Post-process: remove filler words per DCMP §6.01
                captions_vtt = VTTEditor.clean_disfluencies(ai_result["captions_vtt"], source_language)

                # Align caption timings with Whisper word-level timestamps (Bug 5)
                captions_vtt = await _align_captions_with_whisper(captions_vtt, temp_path, job_id)
                # Fix overlapping cues that Gemini occasionally produces
                captions_vtt = VTTEditor.fix_overlapping_cues(captions_vtt)
                ai_result["captions_vtt"] = captions_vtt

                # Fix overlapping cues in AD VTT as well
                ai_result["audio_description_vtt"] = VTTEditor.fix_overlapping_cues(
                    ai_result["audio_description_vtt"]
                )

                # Upload VTT files to GCS using detected language
                captions_gcs_uri = await upload_vtt_to_gcs(
                    ai_result["captions_vtt"],
                    gcs_path(job_doc, source_language, "captions.vtt")
                )

                ad_gcs_uri = await upload_vtt_to_gcs(
                    ai_result["audio_description_vtt"],
                    gcs_path(job_doc, source_language, "ad.vtt")
                )

                # Upload SDH VTT if generated
                sdh_gcs_uri = None
                if sdh_requested and ai_result.get("sdh_captions_vtt"):
                    sdh_gcs_uri = await upload_vtt_to_gcs(
                        ai_result["sdh_captions_vtt"],
                        gcs_path(job_doc, source_language, "sdh_captions.vtt")
                    )

                # Generate descriptive transcript (WCAG 2.1 1.2.1)
                transcript_gcs_uri = None
                try:
                    from ..services.descriptive_transcript import (
                        generate_descriptive_transcript,
                    )
                    transcript_text = generate_descriptive_transcript(
                        ai_result["captions_vtt"],
                        ai_result["audio_description_vtt"]
                    )
                    if transcript_text:
                        transcript_gcs_uri = await upload_vtt_to_gcs(
                            transcript_text,
                            gcs_path(job_doc, source_language, "descriptive_transcript.txt")
                        )
                        logger.info(f"Generated descriptive transcript for job {job_id}, language {source_language}")
                except Exception as e:
                    logger.warning(f"Failed to generate descriptive transcript for job {job_id}: {e}")

                source_lang_output = {
                    "captions_vtt_gcs": captions_gcs_uri,
                    "ad_vtt_gcs": ad_gcs_uri,
                }
                if sdh_gcs_uri:
                    source_lang_output["sdh_captions_vtt_gcs"] = sdh_gcs_uri
                if transcript_gcs_uri:
                    source_lang_output["descriptive_transcript_gcs"] = transcript_gcs_uri

                # Update job with AI results and move to QC for EN approval.
                # Translation pipeline is triggered only after EN QC is approved.
                await db.jobs.update_one(
                    {"_id": job_id},
                    {
                        "$set": {
                            "status": JobStatus.PENDING_QC.value,
                            "source.language": source_language,
                            "source.detected_language": detected_language,
                            "ai.ingestion_json": ai_result,
                            "ai.confidence": ai_result["confidence"],
                            f"outputs.{source_language}": source_lang_output,
                            "updated_at": datetime.utcnow()
                        },
                        "$push": {
                            "review.history": {
                                "at": datetime.utcnow(),
                                "status": JobStatus.PENDING_QC.value,
                                "by": "system"
                            }
                        }
                    }
                )

                broadcast_status_update(
                    job_id,
                    JobStatus.PENDING_QC.value,
                    job_title=job_title,
                    message=f"{job_title} AI processing complete — awaiting EN approval before translation"
                )

                logger.info(f"AI processing complete for job {job_id}, moved to PENDING_QC for EN review")

            finally:
                # Clean up temp file
                os.unlink(temp_path)

    except Exception as e:
        logger.error(f"Ingestion and AI processing failed for job {job_id}: {e}")

        now = datetime.utcnow()
        # Determine step based on current job status
        current_job = await db.jobs.find_one({"_id": job_id}, {"status": 1})
        current_status = current_job.get("status") if current_job else None
        step = "ingestion" if current_status == JobStatus.INGESTING.value else "ai_processing"

        await db.jobs.update_one(
            {"_id": job_id},
            {
                "$set": {
                    "status": JobStatus.PROCESSING_FAILED.value,
                    "failure": {
                        "step": step,
                        "type": type(e).__name__,
                        "message": str(e)[:500],
                        "retriable": True,
                        "occurred_at": now,
                        "retry_count": 0,
                    },
                    "error": {
                        "type": "ingestion_failure",
                        "message": str(e),
                        "timestamp": now.isoformat(),
                    },
                    "updated_at": now,
                }
            }
        )

        raise

    finally:
        client.close()


async def _get_video_duration(video_path: str) -> float:
    """Get video duration using ffprobe"""
    try:
        probe = ffmpeg.probe(video_path)
        duration = float(probe['streams'][0]['duration'])
        return duration
    except Exception as e:
        logger.warning(f"Could not determine video duration: {e}")
        return 0.0


async def _align_captions_with_whisper(captions_vtt: str, video_path: str, job_id: str) -> str:
    """Align caption VTT timings with Whisper word timestamps. Returns original VTT on failure."""
    audio_path = video_path.replace(".mp4", "_captions_align.mp3")
    try:
        # Extract audio at 16kHz mono (optimal for Whisper)
        def _extract():
            result = subprocess.run(
                ["ffmpeg", "-y", "-i", video_path, "-vn", "-acodec", "libmp3lame",
                 "-ar", "16000", "-ac", "1", "-q:a", "5", audio_path],
                capture_output=True, text=True
            )
            if result.returncode != 0:
                raise RuntimeError(f"FFmpeg failed: {result.stderr}")

        await asyncio.to_thread(_extract)

        task_result = transcribe_video_audio_task.apply_async(
            args=[job_id, audio_path], queue="whisper"
        )
        poll_count = 0
        while not task_result.ready():
            await asyncio.sleep(1.0)
            poll_count += 1
            if poll_count > 600:
                logger.warning(f"Whisper timeout for job {job_id}, skipping alignment")
                return captions_vtt

        with allow_join_result():
            result_data = task_result.get(timeout=10)

        words = [
            WordTimestamp(word=w["word"], start=w["start"], end=w["end"])
            for w in result_data.get("words", [])
        ]
        return caption_aligner.align(captions_vtt, words)

    except Exception as e:
        logger.warning(f"Whisper caption alignment failed for job {job_id}: {e} — using Gemini timestamps")
        return captions_vtt
    finally:
        if os.path.exists(audio_path):
            os.unlink(audio_path)
