import asyncio
import logging

from sqlmodel.ext.asyncio.session import AsyncSession

from ..config import Settings
from ..config import get_settings
from ..db.models import Job
from ..db.models import JobSourceStatus
from ..db.models import JobStatus
from ..db.models import Source
from ..errors import AppError
from ..errors import ErrorCategory
from ..errors import classify_unexpected_error
from ..errors import format_error_detail
from ..providers import TranscriptionResult
from . import ServiceBundle
from .transcription import DEFAULT_PROMPT_FILE
from .transcription import transcribe_document_image

logger = logging.getLogger(__name__)


async def advance_job(
    job: Job,
    services: ServiceBundle,
    settings: Settings | None = None,
    session: AsyncSession | None = None,
) -> Job | None:
    """Advance a single job by lifecycle status."""
    settings = settings or get_settings()
    match job.status:
        case JobStatus.QUEUED:
            return await process_queued_job(job=job, services=services, settings=settings, session=session)
        case JobStatus.FAILED:
            if job.retry_count < settings.worker_max_retries:
                return await services.jobs.update_job_state(
                    job_id=job.id,
                    status=JobStatus.QUEUED,
                    retry_count_increment=1,
                    session=session,
                )
            else:
                logger.error(f"Job {job.id} has failed and reached max retries.")
                return
        case _:
            return


async def process_queued_job(
    *,
    job: Job,
    services: ServiceBundle,
    settings: Settings | None = None,
    session: AsyncSession | None = None,
) -> Job | None:
    """Process one complete transcription attempt for a queued job."""
    runtime_settings = settings or get_settings()
    if job.status != JobStatus.QUEUED:
        logger.warning(f"Job {job.id} is not queued. Current status: {job.status}")
        return

    # Transaction A: claim job for processing.
    if session is None:
        job = await services.jobs.mark_job_status(job.id, JobStatus.PROCESSING)
    else:
        # If we're sharing the session, need to make sure setting the Job to PROCESSING is committed before we start
        # the transcription, otherwise other workers may see the job as still QUEUED and try to process it.
        job = await services.jobs.mark_job_status(job.id, JobStatus.PROCESSING, session=session)
        await session.commit()

    source_job = await services.jobs.read_job(job_id=job.id, session=session)
    sources = _resolve_job_sources(source_job)
    if not sources and not source_job.job_sources:
        candidate_sources = await services.transcriptions.list_sources(document_id=job.document_id, session=session)
        sources = list(sorted(candidate_sources, key=lambda item: (item.page_number, item.upload_name.casefold())))

    if not sources:
        return await services.jobs.mark_job_status(job.id, JobStatus.TRANSCRIBED, session=session)

    successful_pages: list[tuple[Source, TranscriptionResult]] = []
    failed_pages: list[tuple[Source, AppError]] = []
    externally_stopped = False

    for source in sources:
        if await _job_no_longer_processing(job_id=job.id, services=services, session=session):
            externally_stopped = True
            break

        started_at = asyncio.get_running_loop().time()
        try:
            result = await asyncio.wait_for(
                transcribe_document_image(source.file_path),
                timeout=runtime_settings.worker_provider_timeout_seconds,
            )
            elapsed_seconds = asyncio.get_running_loop().time() - started_at
            logger.info(
                "Provider response diagnostics operation=worker.provider_response "
                "job_id=%s document_id=%s source_id=%s provider=%s model=%s "
                "finish_reason=%s usage_input_tokens=%s usage_output_tokens=%s usage_total_tokens=%s "
                "latency_seconds=%.3f text_chars=%s text_lines=%s",
                job.id,
                job.document_id,
                source.id,
                result.provider,
                result.model,
                result.finish_reason or "unknown",
                result.usage_input_tokens,
                result.usage_output_tokens,
                result.usage_total_tokens,
                elapsed_seconds,
                len(result.text),
                _line_count(result.text),
            )

            _validate_transcription_quality(result=result, settings=runtime_settings)
            successful_pages.append((source, result))
        except TimeoutError:
            error = AppError(
                f"Provider call timed out after {runtime_settings.worker_provider_timeout_seconds:.1f}s",
                category=ErrorCategory.EXTERNAL_PROVIDER,
                suggestion="Retry the job. If this repeats, verify provider latency and request payload size.",
                retriable=True,
            )
            failed_pages.append((source, error))
            logger.error(
                "Source failed operation=worker.process_job job_id=%s document_id=%s source_id=%s error_id=%s category=%s",
                job.id,
                job.document_id,
                source.id,
                error.error_id,
                error.category.value,
            )
        except Exception as exc:  # noqa: BLE001
            match exc:
                case AppError() as error:
                    pass
                case _:
                    error = classify_unexpected_error(exc, operation="worker.process_job")

            failed_pages.append((source, error))
            logger.error(
                "Source failed operation=worker.process_job job_id=%s document_id=%s source_id=%s error_id=%s category=%s",
                job.id,
                job.document_id,
                source.id,
                error.error_id,
                error.category.value,
            )

        if await _job_no_longer_processing(job_id=job.id, services=services, session=session):
            externally_stopped = True
            break

    terminal_status = JobStatus.TRANSCRIBED
    if externally_stopped:
        terminal_status = JobStatus.FAILED
    elif failed_pages and successful_pages:
        terminal_status = JobStatus.PARTIAL_SUCCESS
    elif failed_pages and not successful_pages:
        terminal_status = JobStatus.FAILED

    updated_job = await _finalize_batch_outcome(
        job=job,
        services=services,
        successful_pages=successful_pages,
        failed_pages=failed_pages,
        status=terminal_status,
        session=session,
    )

    logger.info(
        "Job finished operation=worker.process_job job_id=%s document_id=%s status=%s success_pages=%s failed_pages=%s",
        updated_job.id,
        updated_job.document_id,
        updated_job.status.value,
        len(successful_pages),
        len(failed_pages),
    )
    return updated_job


async def process_next_queued_job(
    *,
    services: ServiceBundle,
    settings: Settings | None = None,
    session: AsyncSession | None = None,
) -> bool:
    """Process the next queued job if one exists."""
    job = await services.jobs.read_next_queued_job(session=session)

    if job is None:
        return False

    await advance_job(job=job, services=services, settings=settings, session=session)
    return True


async def _finalize_transcribed(
    *,
    job: Job,
    services: ServiceBundle,
    result: TranscriptionResult,
    session: AsyncSession | None = None,
) -> Job:
    """Transaction B: job transcription output + TRANSCRIBED in one commit."""
    if session is None:
        async with services.jobs._session_scope() as local_session:
            await services.transcriptions.update_job_transcription(
                job_id=job.id,
                text=result.text,
                error_detail=None,
                provider=result.provider,
                model=result.model,
                prompt_name=result.prompt_name,
                session=local_session,
            )
            updated_job = await services.jobs.mark_job_status(
                job.id,
                JobStatus.TRANSCRIBED,
                session=local_session,
            )
            await local_session.commit()
            return updated_job

    await services.transcriptions.update_job_transcription(
        job_id=job.id,
        text=result.text,
        error_detail=None,
        provider=result.provider,
        model=result.model,
        prompt_name=result.prompt_name,
        session=session,
    )
    updated_job = await services.jobs.mark_job_status(
        job.id,
        JobStatus.TRANSCRIBED,
        session=session,
    )
    await session.commit()
    return updated_job


async def _finalize_retry(
    *,
    job: Job,
    services: ServiceBundle,
    error: AppError,
    settings: Settings,
    session: AsyncSession | None = None,
) -> Job:
    """Transaction C: job error detail + QUEUED + retry increment in one commit."""
    if session is None:
        async with services.jobs._session_scope() as local_session:
            await services.transcriptions.update_job_transcription(
                job_id=job.id,
                text=None,
                error_detail=format_error_detail(error),
                prompt_name=DEFAULT_PROMPT_FILE,
                session=local_session,
            )
            updated_job = await services.jobs.update_job_state(
                job_id=job.id,
                status=JobStatus.QUEUED,
                retry_count_increment=1,
                session=local_session,
            )
            await local_session.commit()
    else:
        await services.transcriptions.update_job_transcription(
            job_id=job.id,
            text=None,
            error_detail=format_error_detail(error),
            prompt_name=DEFAULT_PROMPT_FILE,
            session=session,
        )
        updated_job = await services.jobs.update_job_state(
            job_id=job.id,
            status=JobStatus.QUEUED,
            retry_count_increment=1,
            session=session,
        )
        await session.commit()

    if settings.worker_retry_backoff_seconds > 0:
        await asyncio.sleep(settings.worker_retry_backoff_seconds)
    return updated_job


async def _finalize_failed(
    *,
    job: Job,
    services: ServiceBundle,
    error: AppError,
    session: AsyncSession | None = None,
) -> Job:
    """Transaction B: job error detail + FAILED in one commit."""
    if session is None:
        async with services.jobs._session_scope() as local_session:
            await services.transcriptions.update_job_transcription(
                job_id=job.id,
                text=None,
                error_detail=format_error_detail(error),
                prompt_name=DEFAULT_PROMPT_FILE,
                session=local_session,
            )
            updated_job = await services.jobs.mark_job_status(
                job.id,
                JobStatus.FAILED,
                session=local_session,
            )
            await local_session.commit()
            return updated_job

    await services.transcriptions.update_job_transcription(
        job_id=job.id,
        text=None,
        error_detail=format_error_detail(error),
        prompt_name=DEFAULT_PROMPT_FILE,
        session=session,
    )
    updated_job = await services.jobs.mark_job_status(
        job.id,
        JobStatus.FAILED,
        session=session,
    )
    await session.commit()
    return updated_job


def _resolve_primary_source(job: Job) -> Source | None:
    if not job.job_sources:
        return None
    return next((job_source.source for job_source in job.job_sources if job_source.source is not None), None)


def _resolve_job_sources(job: Job) -> list[Source]:
    """Resolve non-transcribed linked sources for a job in deterministic page order."""
    if not job.job_sources:
        return []

    sources = [
        job_source.source
        for job_source in job.job_sources
        if job_source.source is not None and job_source.status != JobSourceStatus.TRANSCRIBED
    ]
    return list(sorted(sources, key=lambda item: (item.page_number, item.upload_name.casefold())))


async def _job_no_longer_processing(
    *,
    job_id,
    services: ServiceBundle,
    session: AsyncSession | None = None,
) -> bool:
    """Return True when job status changed externally from PROCESSING."""
    latest_job = await services.jobs.read_job(job_id=job_id, session=session)
    return latest_job.status != JobStatus.PROCESSING


async def _finalize_batch_outcome(
    *,
    job: Job,
    services: ServiceBundle,
    successful_pages: list[tuple[Source, TranscriptionResult]],
    failed_pages: list[tuple[Source, AppError]],
    status: JobStatus,
    session: AsyncSession | None = None,
) -> Job:
    """Transaction B: write per-source outcomes and terminal job status atomically."""
    if session is None:
        async with services.jobs._session_scope() as local_session:
            for source, result in successful_pages:
                await services.transcriptions.update_job_source_transcription(
                    job_id=job.id,
                    source_id=source.id,
                    text=result.text,
                    error_detail=None,
                    provider=result.provider,
                    model=result.model,
                    prompt_name=result.prompt_name,
                    session=local_session,
                )

            for source, error in failed_pages:
                await services.transcriptions.update_job_source_transcription(
                    job_id=job.id,
                    source_id=source.id,
                    text=None,
                    error_detail=format_error_detail(error),
                    prompt_name=DEFAULT_PROMPT_FILE,
                    session=local_session,
                )

            updated_job = await services.jobs.mark_job_status(job.id, status, session=local_session)
            await local_session.commit()
            return updated_job

    for source, result in successful_pages:
        await services.transcriptions.update_job_source_transcription(
            job_id=job.id,
            source_id=source.id,
            text=result.text,
            error_detail=None,
            provider=result.provider,
            model=result.model,
            prompt_name=result.prompt_name,
            session=session,
        )

    for source, error in failed_pages:
        await services.transcriptions.update_job_source_transcription(
            job_id=job.id,
            source_id=source.id,
            text=None,
            error_detail=format_error_detail(error),
            prompt_name=DEFAULT_PROMPT_FILE,
            session=session,
        )

    updated_job = await services.jobs.mark_job_status(job.id, status, session=session)
    await session.commit()
    return updated_job


def _validate_transcription_quality(*, result: TranscriptionResult, settings: Settings) -> None:
    text_chars = len(result.text)
    text_lines = _line_count(result.text)

    if settings.worker_fail_on_finish_reason_length and (result.finish_reason or "").lower() == "length":
        raise AppError(
            "Provider output appears truncated (finish_reason=length)",
            category=ErrorCategory.EXTERNAL_PROVIDER,
            suggestion=(
                "Retry the job. If this repeats, use a faster model, reduce input complexity, "
                "or increase provider output budget."
            ),
            retriable=True,
        )

    if settings.worker_min_transcription_chars > 0 and text_chars < settings.worker_min_transcription_chars:
        raise AppError(
            (
                "Transcription output below configured minimum character threshold "
                f"({text_chars} < {settings.worker_min_transcription_chars})"
            ),
            category=ErrorCategory.EXTERNAL_PROVIDER,
            suggestion=(
                "Retry the job. If this repeats, switch model or raise minimum thresholds based on document type."
            ),
            retriable=True,
        )

    if settings.worker_min_transcription_lines > 0 and text_lines < settings.worker_min_transcription_lines:
        raise AppError(
            (
                "Transcription output below configured minimum line threshold "
                f"({text_lines} < {settings.worker_min_transcription_lines})"
            ),
            category=ErrorCategory.EXTERNAL_PROVIDER,
            suggestion=(
                "Retry the job. If this repeats, switch model or raise minimum thresholds based on document type."
            ),
            retriable=True,
        )


def _line_count(text: str) -> int:
    stripped = text.strip()
    if not stripped:
        return 0
    return sum(1 for line in stripped.splitlines() if line.strip())
