"""
Hatchet child workflow: TrackProcessing

Handles individual audio track processing: padding and transcription.
Spawned dynamically by the main diarization pipeline for each track.

Architecture note: This is a separate workflow (not inline tasks in DailyMultitrackPipeline)
because Hatchet workflow DAGs are defined statically, but the number of tracks varies
at runtime. Child workflow spawning via `aio_run()` + `asyncio.gather()` is the
standard pattern for dynamic fan-out. See `process_tracks` in daily_multitrack_pipeline.py.

Note: DB imports (reflector.db.*) are kept inline (deferred) intentionally.
Hatchet workers run in forked processes; fresh imports per task ensure
DB connections are not shared across forks.
"""

from datetime import timedelta

import av
from hatchet_sdk import Context
from pydantic import BaseModel

from reflector.hatchet.client import HatchetClientManager
from reflector.hatchet.constants import TIMEOUT_AUDIO, TIMEOUT_HEAVY
from reflector.hatchet.workflows.models import PadTrackResult, TranscribeTrackResult
from reflector.logger import logger
from reflector.pipelines.transcription_helpers import transcribe_file_with_processor
from reflector.processors.audio_padding_auto import AudioPaddingAutoProcessor
from reflector.storage import get_source_storage, get_transcripts_storage
from reflector.utils.audio_constants import PRESIGNED_URL_EXPIRATION_SECONDS
from reflector.utils.audio_padding import extract_stream_start_time_from_container


class TrackInput(BaseModel):
    """Input for individual track processing."""

    track_index: int
    s3_key: str
    bucket_name: str
    transcript_id: str
    language: str = "en"
    source_platform: str = "daily"
    # Pre-calculated padding in seconds (from filename timestamps for LiveKit).
    # When set, overrides container metadata extraction for start_time.
    padding_seconds: float | None = None


hatchet = HatchetClientManager.get_client()

track_workflow = hatchet.workflow(name="TrackProcessing", input_validator=TrackInput)


@track_workflow.task(
    execution_timeout=timedelta(seconds=TIMEOUT_AUDIO),
    retries=3,
    backoff_factor=2.0,
    backoff_max_seconds=30,
)
async def pad_track(input: TrackInput, ctx: Context) -> PadTrackResult:
    """Pad single audio track with silence for alignment.

    For Daily: extracts stream.start_time from WebM container metadata.
    For LiveKit: uses pre-calculated padding_seconds from filename timestamps
    (OGG files don't have embedded start_time metadata).
    """
    ctx.log(
        f"pad_track: track {input.track_index}, s3_key={input.s3_key}, padding_seconds={input.padding_seconds}"
    )
    logger.info(
        "[Hatchet] pad_track",
        track_index=input.track_index,
        s3_key=input.s3_key,
        transcript_id=input.transcript_id,
        padding_seconds=input.padding_seconds,
    )

    try:
        # Source reads: use platform-specific credentials
        source_storage = get_source_storage(input.source_platform)
        source_url = await source_storage.get_file_url(
            input.s3_key,
            operation="get_object",
            expires_in=PRESIGNED_URL_EXPIRATION_SECONDS,
            bucket=input.bucket_name,
        )

        if input.padding_seconds is not None:
            # Pre-calculated offset (LiveKit: from filename timestamps)
            start_time_seconds = input.padding_seconds
            ctx.log(f"pad_track: using pre-calculated padding={start_time_seconds}s")
        else:
            # Extract from container metadata (Daily: WebM start_time)
            with av.open(source_url) as in_container:
                start_time_seconds = extract_stream_start_time_from_container(
                    in_container, input.track_index, logger=logger
                )

        # If no padding needed, return original S3 key
        if start_time_seconds <= 0:
            logger.info(
                f"Track {input.track_index} requires no padding",
                track_index=input.track_index,
            )
            return PadTrackResult(
                padded_key=input.s3_key,
                bucket_name=input.bucket_name,
                size=0,
                track_index=input.track_index,
            )

        storage_path = f"file_pipeline_hatchet/{input.transcript_id}/tracks/padded_{input.track_index}.webm"

        # Output writes: use transcript storage (our own bucket)
        output_storage = get_transcripts_storage()
        output_url = await output_storage.get_file_url(
            storage_path,
            operation="put_object",
            expires_in=PRESIGNED_URL_EXPIRATION_SECONDS,
        )

        processor = AudioPaddingAutoProcessor()
        result = await processor.pad_track(
            track_url=source_url,
            output_url=output_url,
            start_time_seconds=start_time_seconds,
            track_index=input.track_index,
        )
        file_size = result.size

        ctx.log(f"pad_track complete: track {input.track_index} -> {storage_path}")
        logger.info(
            "[Hatchet] pad_track complete",
            track_index=input.track_index,
            padded_key=storage_path,
        )

        # Return S3 key (not presigned URL) - consumer tasks presign on demand
        # This avoids stale URLs when workflow is replayed
        return PadTrackResult(
            padded_key=storage_path,
            bucket_name=None,  # None = use default transcript storage bucket
            size=file_size,
            track_index=input.track_index,
        )

    except Exception as e:
        logger.error("[Hatchet] pad_track failed", error=str(e), exc_info=True)
        raise


@track_workflow.task(
    parents=[pad_track],
    execution_timeout=timedelta(seconds=TIMEOUT_HEAVY),
    retries=3,
    backoff_factor=2.0,
    backoff_max_seconds=30,
)
async def transcribe_track(input: TrackInput, ctx: Context) -> TranscribeTrackResult:
    """Transcribe audio track using GPU (Modal.com) or local Whisper."""
    ctx.log(f"transcribe_track: track {input.track_index}, language={input.language}")
    logger.info(
        "[Hatchet] transcribe_track",
        track_index=input.track_index,
        language=input.language,
    )

    try:
        pad_result = ctx.task_output(pad_track)
        padded_key = pad_result.padded_key
        bucket_name = pad_result.bucket_name

        if not padded_key:
            raise ValueError("Missing padded_key from pad_track")

        # Presign URL on demand (avoids stale URLs on workflow replay)
        # If bucket_name is set, file is still in the platform's source bucket (no padding applied).
        # If bucket_name is None, padded file was written to our transcript storage.
        if bucket_name:
            storage = get_source_storage(input.source_platform)
        else:
            storage = get_transcripts_storage()

        audio_url = await storage.get_file_url(
            padded_key,
            operation="get_object",
            expires_in=PRESIGNED_URL_EXPIRATION_SECONDS,
            bucket=bucket_name,
        )

        transcript = await transcribe_file_with_processor(audio_url, input.language)

        # Tag all words with speaker index
        for word in transcript.words:
            word.speaker = input.track_index

        ctx.log(
            f"transcribe_track complete: track {input.track_index}, {len(transcript.words)} words"
        )
        logger.info(
            "[Hatchet] transcribe_track complete",
            track_index=input.track_index,
            word_count=len(transcript.words),
        )

        return TranscribeTrackResult(
            words=transcript.words,
            track_index=input.track_index,
        )

    except Exception as e:
        logger.error("[Hatchet] transcribe_track failed", error=str(e), exc_info=True)
        raise
