"""
OpenLineage integration for pycharter ETL pipelines.

Emits START / COMPLETE / FAIL RunEvents with standard and custom facets
derived from pipeline config and PipelineResult.

All openlineage imports are lazy so the module is importable even when
openlineage-python is not installed.  When the package is missing or
lineage is disabled, ``create_lineage_emitter`` returns ``None`` and
the pipeline runs with zero overhead.
"""

from __future__ import annotations

import logging
import re
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from pycharter.etl_generator.config_models import LineageConfig
    from pycharter.etl_generator.result import PipelineResult

logger = logging.getLogger(__name__)

PRODUCER = "https://github.com/optophi/pycharter"
_CREDENTIAL_RE = re.compile(r"://[^@/]*@")

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _sanitize_url(url: str) -> str:
    """Strip credentials from a connection string for safe lineage metadata."""
    return _CREDENTIAL_RE.sub("://***@", url)


# ---------------------------------------------------------------------------
# Dataset mappers
# ---------------------------------------------------------------------------


def dataset_from_extract_config(config: dict[str, Any]) -> tuple[str, str]:
    """Return (namespace, name) for an OpenLineage InputDataset from extract config."""
    cfg_type = (config.get("type") or "").lower()

    if cfg_type == "http":
        ns = config.get("base_url") or config.get("url") or "http"
        name = config.get("api_endpoint") or config.get("endpoint") or "/"
        return ns, name

    if cfg_type == "database":
        db_cfg = config.get("database") or {}
        url = db_cfg.get("url") or config.get("connection_string") or "database"
        name = config.get("query", "query")[:120]
        return _sanitize_url(url), name

    if cfg_type == "file":
        path = config.get("path") or config.get("file_path") or "unknown"
        return "file", path

    if cfg_type == "cloud_storage":
        storage = config.get("storage") or config
        provider = storage.get("provider", "cloud")
        bucket = storage.get("bucket") or storage.get("container") or "unknown"
        path = storage.get("path") or storage.get("prefix") or ""
        return f"{provider}://{bucket}", path

    if cfg_type == "mongodb":
        mongo = config.get("mongodb") or {}
        url = mongo.get("connection_string") or "mongodb"
        collection = mongo.get("collection") or "unknown"
        return _sanitize_url(url), collection

    return "unknown", cfg_type or "unknown"


def dataset_from_load_config(config: dict[str, Any]) -> tuple[str, str]:
    """Return (namespace, name) for an OpenLineage OutputDataset from load config."""
    cfg_type = (config.get("type") or "").lower()

    if cfg_type in ("postgres", "postgresql", "database", "sqlite"):
        db_cfg = config.get("database") or {}
        url = db_cfg.get("url") or config.get("connection_string") or cfg_type
        schema = config.get("schema_name") or "public"
        table = config.get("table") or "unknown"
        return _sanitize_url(url), f"{schema}.{table}"

    if cfg_type == "file":
        path = config.get("path") or config.get("file_path") or "unknown"
        return "file", path

    if cfg_type == "cloud_storage":
        storage = config.get("storage") or config
        provider = storage.get("provider", "cloud")
        bucket = storage.get("bucket") or storage.get("container") or "unknown"
        path = storage.get("path") or storage.get("prefix") or ""
        return f"{provider}://{bucket}", path

    if cfg_type == "mongodb":
        mongo = config.get("mongodb") or {}
        url = mongo.get("connection_string") or "mongodb"
        collection = mongo.get("collection") or "unknown"
        return _sanitize_url(url), collection

    return "unknown", cfg_type or "unknown"


# ---------------------------------------------------------------------------
# Facet builders
# ---------------------------------------------------------------------------


def build_run_facets(result: PipelineResult) -> dict[str, Any]:
    """Build run-level facets from a PipelineResult.

    Includes a custom ``pycharter_pipeline`` facet with pipeline statistics
    and, when errors are present, the standard ``errorMessage`` facet.
    """
    from openlineage.client.facet_v2 import error_message_run as _em

    facets: dict[str, Any] = {}

    # Standard errorMessage facet (only when the run has errors)
    if result.errors:
        facets["errorMessage"] = _em.ErrorMessageRunFacet(
            message=result.errors[0],
            programmingLanguage="python",
            stackTrace="\n".join(result.errors) if len(result.errors) > 1 else None,
        )

    # Custom pycharter_pipeline facet — lightweight dict; serialised as-is.
    facets["pycharter_pipeline"] = {
        "_producer": PRODUCER,
        "_schemaURL": f"{PRODUCER}#/definitions/PycharterPipelineRunFacet",
        "rows_extracted": result.rows_extracted,
        "rows_transformed": result.rows_transformed,
        "rows_loaded": result.rows_loaded,
        "rows_failed": result.rows_failed,
        "rows_quarantined_extract": result.rows_quarantined_extract,
        "rows_quarantined_load": result.rows_quarantined_load,
        "total_quarantined": result.total_quarantined,
        "batches_processed": result.batches_processed,
        "duration_seconds": result.duration_seconds,
        "validation_errors_extract": len(result.validation_errors_extract),
        "validation_errors_load": len(result.validation_errors_load),
    }

    return facets


def build_input_facets(result: PipelineResult) -> dict[str, Any]:
    """Build input-dataset facets (standard dataQualityMetrics)."""
    from openlineage.client.facet_v2 import data_quality_metrics_input_dataset as _dqm

    return {
        "dataQualityMetrics": _dqm.DataQualityMetricsInputDatasetFacet(
            rowCount=result.rows_extracted,
            columnMetrics={},
        ),
    }


def build_output_facets(result: PipelineResult) -> dict[str, Any]:
    """Build output-dataset facets (standard outputStatistics)."""
    from openlineage.client.facet_v2 import output_statistics_output_dataset as _os

    return {
        "outputStatistics": _os.OutputStatisticsOutputDatasetFacet(
            rowCount=result.rows_loaded,
        ),
    }


def run_event_to_dict(event: Any) -> dict[str, Any] | None:
    """Serialize an OpenLineage RunEvent to a JSON-serializable dict.

    Used to capture emitted events for API responses (e.g. ETL UI).
    Returns None if serialization fails.
    """
    try:
        import json

        from openlineage.client.serde import Serde

        return json.loads(Serde.to_json(event))
    except Exception:
        return None


# ---------------------------------------------------------------------------
# LineageEmitter
# ---------------------------------------------------------------------------


class LineageEmitter:
    """Thin wrapper around OpenLineageClient that emits pipeline RunEvents."""

    def __init__(self, config: LineageConfig) -> None:
        from openlineage.client import OpenLineageClient

        transport_config: dict[str, Any] | None = None
        if config.backend_url:
            transport_config = {
                "transport": {"type": "http", "url": config.backend_url}
            }

        self._client = OpenLineageClient(config=transport_config)
        self._namespace = config.namespace

    def start(
        self,
        job_name: str,
        run_id: str,
        inputs: list[tuple[str, str]],
    ) -> dict[str, Any] | None:
        """Emit a START event (before extraction begins). Returns payload dict for UI."""
        from openlineage.client.event_v2 import (
            InputDataset,
            Job,
            Run,
            RunEvent,
            RunState,
        )

        event = RunEvent(
            eventType=RunState.START,
            eventTime=datetime.now(timezone.utc).isoformat(),
            run=Run(runId=run_id),
            job=Job(namespace=self._namespace, name=job_name),
            inputs=[InputDataset(ns, name) for ns, name in inputs],
            outputs=[],
            producer=PRODUCER,
        )
        self._client.emit(event)
        return run_event_to_dict(event)

    def complete(
        self,
        job_name: str,
        run_id: str,
        inputs: list[tuple[str, str]],
        outputs: list[tuple[str, str]],
        result: PipelineResult,
    ) -> dict[str, Any] | None:
        """Emit a COMPLETE event with run, input, and output facets. Returns payload dict."""
        from openlineage.client.event_v2 import (
            InputDataset,
            Job,
            OutputDataset,
            Run,
            RunEvent,
            RunState,
        )

        run_facets = build_run_facets(result)
        in_facets = build_input_facets(result)
        out_facets = build_output_facets(result)

        event = RunEvent(
            eventType=RunState.COMPLETE,
            eventTime=datetime.now(timezone.utc).isoformat(),
            run=Run(runId=run_id, facets=run_facets),
            job=Job(namespace=self._namespace, name=job_name),
            inputs=[
                InputDataset(ns, name, inputFacets=in_facets) for ns, name in inputs
            ],
            outputs=[
                OutputDataset(ns, name, outputFacets=out_facets) for ns, name in outputs
            ],
            producer=PRODUCER,
        )
        self._client.emit(event)
        return run_event_to_dict(event)

    def fail(
        self,
        job_name: str,
        run_id: str,
        inputs: list[tuple[str, str]],
        result: PipelineResult,
    ) -> dict[str, Any] | None:
        """Emit a FAIL event with error details and partial stats. Returns payload dict."""
        from openlineage.client.event_v2 import (
            InputDataset,
            Job,
            Run,
            RunEvent,
            RunState,
        )

        run_facets = build_run_facets(result)

        event = RunEvent(
            eventType=RunState.FAIL,
            eventTime=datetime.now(timezone.utc).isoformat(),
            run=Run(runId=run_id, facets=run_facets),
            job=Job(namespace=self._namespace, name=job_name),
            inputs=[InputDataset(ns, name) for ns, name in inputs],
            outputs=[],
            producer=PRODUCER,
        )
        self._client.emit(event)
        return run_event_to_dict(event)


# ---------------------------------------------------------------------------
# Factory
# ---------------------------------------------------------------------------


def create_lineage_emitter(settings: dict[str, Any]) -> LineageEmitter | None:
    """Create a LineageEmitter from pipeline settings, or return None.

    Returns None when:
    - lineage is not configured or ``enabled`` is False
    - openlineage-python is not installed
    """
    lineage_raw = settings.get("lineage")
    if not lineage_raw:
        return None

    if isinstance(lineage_raw, dict):
        if not lineage_raw.get("enabled", False):
            return None
        from pycharter.etl_generator.config_models import LineageConfig

        config = LineageConfig(**lineage_raw)
    else:
        # Already a LineageConfig instance (e.g. from parsed settings)
        if not getattr(lineage_raw, "enabled", False):
            return None
        config = lineage_raw

    try:
        return LineageEmitter(config)
    except ImportError:
        logger.debug(
            "openlineage-python not installed; lineage emission disabled. "
            "Install with: pip install pycharter[lineage]"
        )
        return None
    except Exception:
        logger.warning("Failed to create lineage emitter", exc_info=True)
        return None
