"""Classe de base pour les pipelines spécialisés CEZAM."""

from __future__ import annotations

import json
import logging
import time
from abc import ABC
from typing import TYPE_CHECKING

from opentelemetry import metrics, trace
from opentelemetry.trace import Status, StatusCode

from cezam_lib.cezam_shared import datalake_paths
from cezam_lib.cezam_shared.exceptions import MinIOError, NonRetryableError, RetryableError
from .extractor import DataExtractor
from .messages import FusionMessage, PipelineMessage

if TYPE_CHECKING:
    from cezam_lib.cezam_shared.datalake_client import DatalakeClient
    from cezam_lib.cezam_shared.rabbitmq import RabbitMQConsumer, RabbitMQPublisher


class _JsonFormatter(logging.Formatter):
    """Formateur de logs JSON avec corrélation trace/simulation."""

    def format(self, record: logging.LogRecord) -> str:
        span = trace.get_current_span()
        ctx = span.get_span_context()
        payload: dict = {
            "timestamp": self.formatTime(record, self.datefmt),
            "level": record.levelname,
            "logger": record.name,
            "message": record.getMessage(),
            "pipeline": getattr(record, "pipeline", None),
            "simulation_id": getattr(record, "simulation_id", None),
            "duration_ms": getattr(record, "duration_ms", None),
            "trace_id": format(ctx.trace_id, "032x") if ctx.is_valid else None,
            "span_id": format(ctx.span_id, "016x") if ctx.is_valid else None,
        }
        if record.exc_info:
            payload["error"] = self.formatException(record.exc_info)
        return json.dumps(payload, ensure_ascii=False)


def _setup_json_logging() -> None:
    handler = logging.StreamHandler()
    handler.setFormatter(_JsonFormatter())
    root = logging.getLogger()
    root.handlers = [handler]
    root.setLevel(logging.INFO)


class BasePipeline(ABC):
    """Classe de base pour les pipelines spécialisés.

    Gère le flux complet de traitement d'un message:
    1. Parse le PipelineMessage
    2. Lit l'OCR depuis le datalake
    3. Appelle l'extracteur
    4. Écrit le résultat sur le datalake
    5. Publie vers la queue fusion
    """

    FUSION_QUEUE = "fusion"
    FUSION_EXCHANGE = ""

    def __init__(
        self,
        datalake_client: DatalakeClient,
        publisher: RabbitMQPublisher,
        consumer: RabbitMQConsumer,
        extractor: DataExtractor,
        queue_name: str,
        pipeline_name: str,
    ) -> None:
        """Initialise le pipeline.

        Args:
            datalake_client: Client datalake pour lecture/écriture JSON.
            publisher: Publisher RabbitMQ pour envoyer vers fusion.
            consumer: Consumer RabbitMQ pour recevoir les messages.
            extractor: Extracteur de données spécialisé.
            queue_name: Nom de la queue à consommer.
            pipeline_name: Nom du pipeline (pour les chemins datalake).
        """
        _setup_json_logging()
        self._datalake = datalake_client
        self._publisher = publisher
        self._consumer = consumer
        self._extractor = extractor
        self._queue_name = queue_name
        self._pipeline_name = pipeline_name
        self._tracer = trace.get_tracer(pipeline_name)
        self._logger = logging.getLogger(pipeline_name)

        meter = metrics.get_meter(pipeline_name)
        self._counter_processed = meter.create_counter(
            name="pipeline_documents_processed_total",
            description="Nombre de documents traités par ce pipeline",
            unit="1",
        )
        self._counter_failed = meter.create_counter(
            name="pipeline_documents_failed_total",
            description="Nombre de documents en erreur",
            unit="1",
        )
        self._histogram_duration = meter.create_histogram(
            name="pipeline_processing_duration_seconds",
            description="Durée de traitement d'un document de bout en bout",
            unit="s",
        )
        self._histogram_quality = meter.create_histogram(
            name="pipeline_result_quality_score",
            description="Score de confiance moyen du résultat extrait (0.0 à 1.0)",
            unit="1",
        )

    def process_message(self, message: dict) -> None:
        """Traite un message de la queue.

        Args:
            message: Message brut (dict) reçu de RabbitMQ.

        Raises:
            RetryableError: Si une erreur temporaire survient.
            NonRetryableError: Si une erreur définitive survient.
        """
        start_time = time.perf_counter()
        labels = {"pipeline": self._pipeline_name}
        pipeline_msg = None

        with self._tracer.start_as_current_span("process_message") as span:
            try:
                pipeline_msg = self._parse_message(message, span)
                self._logger.info(
                    "Début traitement",
                    extra={"pipeline": self._pipeline_name, "simulation_id": pipeline_msg.simulation_id},
                )
                ocr_data = self._read_ocr(pipeline_msg.ocr_json_path, span)
                result = self._extract_data(ocr_data, span)
                output_path = self._write_result(pipeline_msg, result, span)
                duration_ms = int((time.perf_counter() - start_time) * 1000)
                self._publish_fusion(pipeline_msg, output_path, result, duration_ms, span)

                duration_s = duration_ms / 1000.0
                self._counter_processed.add(1, {**labels, "status": "success"})
                self._histogram_duration.record(duration_s, labels)
                self._logger.info(
                    "Document traité avec succès",
                    extra={"pipeline": self._pipeline_name, "simulation_id": pipeline_msg.simulation_id, "duration_ms": duration_ms},
                )

            except (RetryableError, NonRetryableError):
                self._counter_processed.add(1, {**labels, "status": "error"})
                self._counter_failed.add(1, labels)
                sim_id = pipeline_msg.simulation_id if pipeline_msg else None
                self._logger.error(
                    "Erreur de traitement du document",
                    extra={"pipeline": self._pipeline_name, "simulation_id": sim_id},
                    exc_info=True,
                )
                raise

            except Exception as e:
                self._counter_processed.add(1, {**labels, "status": "error"})
                self._counter_failed.add(1, labels)
                span.set_status(Status(StatusCode.ERROR, str(e)))
                span.record_exception(e)
                sim_id = pipeline_msg.simulation_id if pipeline_msg else None
                self._logger.error(
                    "Erreur inattendue",
                    extra={"pipeline": self._pipeline_name, "simulation_id": sim_id},
                    exc_info=True,
                )
                raise NonRetryableError(f"Erreur inattendue: {e}") from e

    def _parse_message(self, message: dict, span: trace.Span) -> PipelineMessage:
        """Parse et valide le message entrant.

        Args:
            message: Message brut.
            span: Span OTel parent.

        Returns:
            PipelineMessage validé.

        Raises:
            NonRetryableError: Si le message est invalide.
        """
        with self._tracer.start_as_current_span("parse_message"):
            try:
                return PipelineMessage.model_validate(message)
            except Exception as e:
                raise NonRetryableError(f"Message invalide: {e}") from e

    def _read_ocr(self, ocr_path: str, span: trace.Span) -> dict:
        """Lit les données OCR depuis le datalake.

        Args:
            ocr_path: Chemin du fichier OCR (sans env_prefix).
            span: Span OTel parent.

        Returns:
            Données OCR (dict).

        Raises:
            RetryableError: Si le datalake est temporairement indisponible.
            NonRetryableError: Si le fichier n'existe pas.
        """
        with self._tracer.start_as_current_span("read_ocr") as read_span:
            read_span.set_attribute("ocr_path", ocr_path)
            try:
                return self._datalake.get_json(ocr_path)
            except MinIOError as e:
                error_msg = str(e)
                if "inexistant" in error_msg.lower() or "nosuchkey" in error_msg.lower():
                    raise NonRetryableError(f"Fichier OCR introuvable: {ocr_path}") from e
                raise RetryableError(f"Erreur datalake temporaire: {e}") from e

    def _extract_data(self, ocr_data: dict, span: trace.Span) -> dict:
        """Extrait les données structurées depuis l'OCR.

        Args:
            ocr_data: Données OCR.
            span: Span OTel parent.

        Returns:
            Données extraites (dict).

        Raises:
            RetryableError: Si l'extraction peut être retentée.
            NonRetryableError: Si l'extraction a échoué définitivement.
        """
        with self._tracer.start_as_current_span("extract_data"):
            result = self._extractor.extract(ocr_data)
            return result.model_dump()

    def _write_result(
        self, pipeline_msg: PipelineMessage, result: dict, span: trace.Span
    ) -> str:
        """Écrit le résultat sur le datalake.

        Args:
            pipeline_msg: Message pipeline original.
            result: Données extraites.
            span: Span OTel parent.

        Returns:
            Chemin du fichier écrit sur le datalake (sans env_prefix).

        Raises:
            RetryableError: Si le datalake est temporairement indisponible.
        """
        output_path = datalake_paths.pipeline_result_path(
            pipeline_msg.simulation_id, self._pipeline_name, pipeline_msg.doc_name
        )

        with self._tracer.start_as_current_span("write_result") as write_span:
            write_span.set_attribute("output_path", output_path)
            try:
                self._datalake.put_json(output_path, result)
                return output_path
            except MinIOError as e:
                raise RetryableError(f"Erreur datalake lors de l'écriture: {e}") from e

    def _publish_fusion(
        self,
        pipeline_msg: PipelineMessage,
        output_path: str,
        result: dict,
        duration_ms: int,
        span: trace.Span,
    ) -> None:
        """Publie le message vers la queue fusion.

        Extrait les métriques réelles du résultat (quality, action,
        confidence, field_count) pour que le service fusion puisse
        prendre une décision sans relire le JSON MinIO.

        Args:
            pipeline_msg: Message pipeline original.
            output_path: Chemin du résultat sur le datalake.
            result: Données extraites.
            duration_ms: Durée du traitement en ms.
            span: Span OTel parent.

        Raises:
            RetryableError: Si RabbitMQ est temporairement indisponible.
        """
        meta = _extract_result_meta(result)

        fusion_msg = FusionMessage(
            simulation_id=pipeline_msg.simulation_id,
            simulation_public_id=pipeline_msg.simulation_public_id,
            document_type=pipeline_msg.document_type,
            pipeline_output_path=output_path,
            status=meta["status"],
            quality=meta["quality"],
            action=meta["action"],
            fields_extracted=meta["fields_extracted"],
            confidence_avg=meta["confidence_avg"],
            warnings_count=meta["warnings_count"],
            source_documents=[pipeline_msg.doc_name],
            partner=meta["partner"],
            pipeline_name=self._pipeline_name,
            duration_ms=duration_ms,
            correlation_id=pipeline_msg.correlation_id,
            n_doc=pipeline_msg.n_doc,
            n_doc_total=pipeline_msg.n_doc_total,
        )

        with self._tracer.start_as_current_span("publish_fusion") as pub_span:
            pub_span.set_attribute("fusion_queue", self.FUSION_QUEUE)
            pub_span.set_attribute("result_quality", meta["quality"])
            pub_span.set_attribute("result_action", meta["action"])
            pub_span.set_attribute("confidence_avg", meta["confidence_avg"])
            try:
                self._publisher.publish(
                    exchange=self.FUSION_EXCHANGE,
                    routing_key=self.FUSION_QUEUE,
                    message=fusion_msg,
                )
                self._histogram_quality.record(
                    meta["confidence_avg"],
                    {"pipeline": self._pipeline_name},
                )
            except Exception as e:
                raise RetryableError(f"Erreur publication vers fusion: {e}") from e

    def run(self) -> None:
        """Lance le consumer pour traiter les messages."""
        self._consumer.consume(
            queue=self._queue_name,
            callback=self.process_message,
        )


def _extract_result_meta(result: dict) -> dict:
    """Extrait les métriques unifiées depuis le résultat d'un pipeline.

    Gère les deux formats :
      - DDP : {status, confidence, quality, action, field_count, warnings, partner}
      - Identité : {status, persons: [{global_confidence, quality, action, ...}], person_count}
    """
    status_raw = (result.get("status") or "ERROR").upper()

    # Confiance, quality, action — format DDP (champs racine)
    confidence = result.get("confidence", 0.0)
    quality = result.get("quality", "POOR")
    action = result.get("action", "REJECT")
    fields = result.get("field_count", 0)
    warnings = result.get("warnings", [])

    # Format identité — agréger depuis persons[]
    persons = result.get("persons", [])
    if persons and not result.get("confidence"):
        confidences = [p.get("global_confidence", 0.0) for p in persons]
        confidence = sum(confidences) / len(confidences) if confidences else 0.0

        # Pire quality/action parmi les personnes
        quality_order = {"EXCELLENT": 0, "GOOD": 1, "ACCEPTABLE": 2, "POOR": 3}
        action_order = {"AUTO_VALIDATE": 0, "MANUAL_REVIEW": 1, "REJECT": 2}
        qualities = [p.get("quality", "POOR") for p in persons]
        actions = [p.get("action", "REJECT") for p in persons]
        quality = max(qualities, key=lambda q: quality_order.get(q, 3))
        action = max(actions, key=lambda a: action_order.get(a, 2))

        fields = sum(len(p.get("extracted_data", {})) for p in persons)
        warnings = []
        for p in persons:
            warnings.extend(p.get("warnings", []))

    # Mapper status vers success/partial/error
    if status_raw == "SUCCESS":
        status = "success"
    elif status_raw == "PARTIAL":
        status = "partial"
    else:
        status = "error"

    return {
        "status": status,
        "quality": quality,
        "action": action,
        "confidence_avg": min(max(confidence, 0.0), 1.0),
        "fields_extracted": fields,
        "warnings_count": len(warnings),
        "partner": result.get("partner"),
    }
