"""Tareas ARQ para procesamiento paralelo de documentos.

Cada tarea recibe la ruta de un archivo temporal en el volumen compartido
`/tmp/ips_queue/`, lo procesa y actualiza la tarea en la base de datos.
El archivo se elimina al finalizar (éxito o error).
"""
from __future__ import annotations

import asyncio
import logging
import os
from typing import Any

from sqlalchemy.orm.attributes import flag_modified
from infrastructure.database import ExtractionRow, Task, session_scope, utcnow
from infrastructure.serialization import flatten_extraction_result
from api.desktop_ingest import (
    _finalize_radicado_if_ready,
    _paths_for_type,
    _service_for_model,
    mark_task_completed_if_done,
)
from domain.profiles import get_profile
from domain.prompts import PromptBuilder
from domain.normalizer import normalize_to_schema
from domain.rate_limiter import get_rate_limiter
from domain.config import settings
from domain.usage_tracking import record_api_usage

logger = logging.getLogger(__name__)


def _service_for_request(
    ctx: dict[str, Any], request: dict[str, Any]
) -> Any:
    """Resuelve y reutiliza el servicio correspondiente al modelo solicitado."""
    base_service = ctx["service"]
    requested_model = str(request.get("model") or "").strip() or None
    if not requested_model:
        return base_service

    current_model = getattr(getattr(base_service, "extractor", None), "model", None)
    if current_model == requested_model:
        return base_service

    services_by_model = ctx.setdefault("services_by_model", {})
    service = services_by_model.get(requested_model)
    if service is None:
        service = _service_for_model(base_service, requested_model)
        services_by_model[requested_model] = service
    return service


async def process_document(
    ctx: dict[str, Any],
    *,
    task_id: str,
    batch_id: str,
    queue_path: str,
    radicado: str,
    filename: str,
    doc_type: str | None,
    instruction: str | None,
    tiff_page_range: list[int] | None,
) -> dict[str, Any]:
    """Extrae datos de un documento y los persiste en la BD.

    Args:
        ctx: Contexto ARQ con 'service' (InvoiceProcessingService).
        task_id: ID de la tarea parent.
        batch_id: ID del lote de extracción.
        queue_path: Ruta absoluta al archivo en el volumen compartido.
        radicado: Radicado al que pertenece el documento.
        filename: Nombre original del archivo.
        doc_type: Tipo documental (None en modo carpeta general).
        instruction: Instrucción de extracción personalizada.
        tiff_page_range: Rango de páginas TIFF [inicio, fin].
    """
    rate_limiter = get_rate_limiter()
    tiff_range = tuple(tiff_page_range) if tiff_page_range else None

    logger.info("Procesando %s / %s (task=%s)", radicado, filename, task_id)

    try:
        # Obtener config de la tarea desde BD
        def _load_task_config() -> dict[str, Any] | None:
            with session_scope() as session:
                task = session.get(Task, task_id)
                if task is None:
                    return None
                return dict(task.request or {})

        request = await asyncio.to_thread(_load_task_config)
        if request is None:
            logger.warning("Tarea %s no encontrada — saltando %s", task_id, filename)
            return {"skipped": True, "reason": "task_not_found"}

        service = _service_for_request(ctx, request)
        effective_model = getattr(getattr(service, "extractor", None), "model", "?")
        logger.info(
            "Modelo resuelto para task=%s: %s",
            task_id,
            effective_model,
        )

        fields: list[str] = list(request.get("fields") or [])
        alias_map: dict[str, str] = dict(request.get("alias_map") or {})
        structure_mode: str = str(request.get("structure_mode") or "radicados")
        doc_types: list = list(request.get("doc_types") or [])
        prof = get_profile("invoice")

        selected_paths = _paths_for_type(
            fields,
            doc_type if structure_mode == "radicados" and doc_types else None,
        )

        if doc_type and structure_mode == "radicados" and doc_types:
            try:
                prompt = PromptBuilder.build_for_type(
                    doc_type, selected_paths, {}, profile_name=prof.get("name", "invoice")
                )
            except Exception:
                prompt = PromptBuilder.build(
                    selected_paths, {}, profile_name=prof.get("name", "invoice")
                )
        else:
            prompt = instruction or PromptBuilder.build(
                selected_paths, {}, profile_name=prof.get("name", "invoice")
            )

        # Esperar turno en el rate limiter antes de llamar a OpenAI
        estimated_tokens = max(len(selected_paths) * 80, 2000)
        await rate_limiter.acquire(estimated_tokens=estimated_tokens)

        # Extracción + normalización (CPU + IO bound, en thread pool)
        def _extract_and_normalize() -> tuple[dict, dict, dict | None, dict | None]:
            data, extraction_usage = service.extract_full_from_pdf(
                queue_path,
                instruction=prompt,
                tiff_page_range=tiff_range,
            )
            raw_data: dict = data if isinstance(data, dict) else {"data": data}
            if isinstance(data, dict):
                data["radicado"] = radicado
            else:
                data = {"radicado": radicado, "data": data}

            norm = normalize_to_schema(
                data,
                paths=selected_paths,
                profile_name=prof.get("name", "invoice"),
            )
            normalized_obj = data
            normalizer_usage = None
            if isinstance(norm, dict):
                normalized_obj = norm.get("normalized") or data
                normalizer_usage = norm.get("_usage")
            return raw_data, normalized_obj, extraction_usage, normalizer_usage

        _raw, normalized_obj, extraction_usage, normalizer_usage = await asyncio.to_thread(_extract_and_normalize)

        # Persistir resultado en BD
        def _persist() -> dict[str, Any]:
            with session_scope() as session:
                # Lock evita lost-update en result.processed con N workers concurrentes.
                task = session.get(Task, task_id, with_for_update=True)
                if task is None:
                    return {"skipped": True, "reason": "task_not_found"}

                result = dict(task.result or {})
                result = record_api_usage(
                    result, phase="extract", model=str(effective_model), usage=extraction_usage,
                    radicado=radicado, filename=filename,
                )
                result = record_api_usage(
                    result, phase="normalize", model=str(settings.api.normalizer_model), usage=normalizer_usage,
                    radicado=radicado, filename=filename,
                )
                task.result = result

                if structure_mode == "radicados" and doc_types and doc_type:
                    partials = dict(result.get("partials") or {})
                    radicado_partials = dict(partials.get(radicado) or {})
                    radicado_partials[doc_type] = normalized_obj or {}
                    partials[radicado] = radicado_partials
                    result["partials"] = partials
                    task.result = result
                    session.flush()
                    finalized = _finalize_radicado_if_ready(
                        session,
                        task=task,
                        batch_id=batch_id,
                        radicado=radicado,
                        request=request,
                        fields=fields,
                        alias_map=alias_map,
                    )
                    mark_task_completed_if_done(session, task, batch_id)
                    return {
                        "radicado": radicado,
                        "filename": filename,
                        "doc_type": doc_type,
                        "finalized": finalized,
                        "processed": int((task.result or {}).get("processed") or 0),
                    }

                fused = {"radicado": radicado, "_source_pdf": filename}
                if doc_type:
                    fused[doc_type] = normalized_obj or {}
                else:
                    fused.update(
                        normalized_obj if isinstance(normalized_obj, dict) else {"data": normalized_obj}
                    )

                row_index = int(result.get("row_count") or 0)
                for row in flatten_extraction_result(fused, fields, alias_map):
                    row_index += 1
                    session.add(
                        ExtractionRow(
                            batch_id=batch_id,
                            radicado=radicado,
                            source_pdf=filename,
                            row_index=row_index,
                            row_json=row,
                            raw_json=fused,
                        )
                    )

                result["row_count"] = row_index
                result["processed"] = int(result.get("processed") or 0) + 1
                task.result = result
                flag_modified(task, "result")
                task.updated_at = utcnow()
                session.flush()
                mark_task_completed_if_done(session, task, batch_id)
                return {
                    "radicado": radicado,
                    "filename": filename,
                    "doc_type": doc_type,
                    "finalized": True,
                    "processed": int(result.get("processed") or 0),
                }

        payload = await asyncio.to_thread(_persist)
        logger.info("Completado %s / %s — payload=%s", radicado, filename, payload)
        return payload

    except Exception as exc:
        logger.exception("Error procesando %s / %s: %s", radicado, filename, exc)
        # Registrar error en la tarea sin fallar el worker
        def _record_error() -> None:
            try:
                with session_scope() as session:
                    task = session.get(Task, task_id, with_for_update=True)
                    if task is None:
                        return
                    result = dict(task.result or {})
                    errors = list(result.get("errors") or [])
                    errors.append({
                        "radicado": radicado,
                        "filename": filename,
                        "error": str(exc),
                    })
                    result["errors"] = errors
                    result["processed"] = int(result.get("processed") or 0) + 1
                    task.result = result
                    flag_modified(task, "result")
                    task.updated_at = utcnow()
                    mark_task_completed_if_done(session, task, batch_id)
            except Exception as inner:
                logger.error("No se pudo registrar el error en BD: %s", inner)

        await asyncio.to_thread(_record_error)
        return {"error": str(exc), "radicado": radicado, "filename": filename}

    finally:
        # Eliminar archivo temporal del volumen compartido
        try:
            if os.path.exists(queue_path):
                os.unlink(queue_path)
        except OSError as e:
            logger.warning("No se pudo eliminar %s: %s", queue_path, e)