# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
"""Registry store comparison operations.

This module compares two registry stores (e.g. a candidate dev store against
the prod reference) and reports differences at both the index and per-connector
artifact level.

Usage::

    from airbyte_ops_mcp.registry.compare import compare_stores

    result = compare_stores(
        store_bucket="dev-airbyte-cloud-connector-metadata-service-2",
        store_prefix="dev/20260306-mirror-compile/",
        reference_bucket="prod-airbyte-cloud-connector-metadata-service",
        reference_prefix="",
        connector_filter=["source-faker"],
    )

High-level algorithm
--------------------
1. List connectors in both stores (``metadata/airbyte/*/latest/``).
2. For connectors present in both, compare per-connector artifacts:
   ``metadata.yaml``, ``cloud.json``, ``oss.json``, ``spec.json``.
3. Compare global registry indexes:
   ``registries/v0/cloud_registry.json``, ``registries/v0/oss_registry.json``.
4. Return a structured result with per-connector and index-level diffs.
"""

from __future__ import annotations

import json
import logging
import os
import sys
from dataclasses import dataclass, field
from typing import Any

import gcsfs

from airbyte_ops_mcp.registry._constants import METADATA_FOLDER

logger = logging.getLogger(__name__)

# Artifact files to compare per connector version
_ARTIFACT_FILES = ("metadata.yaml", "cloud.json", "oss.json", "spec.json")

# Global index files to compare
_INDEX_FILES = (
    "registries/v0/cloud_registry.json",
    "registries/v0/oss_registry.json",
)


# ---------------------------------------------------------------------------
# Result models
# ---------------------------------------------------------------------------


@dataclass
class ArtifactDiff:
    """Difference found in a single artifact file."""

    connector: str
    file: str
    status: str  # "only_in_store", "only_in_reference", "content_differs"
    details: str = ""


@dataclass
class ConnectorDiff:
    """Summary of differences for a single connector."""

    connector: str
    status: str  # "only_in_store", "only_in_reference", "artifacts_differ", "match"
    artifact_diffs: list[ArtifactDiff] = field(default_factory=list)


@dataclass
class IndexDiff:
    """Difference found in a global index file."""

    file: str
    status: str  # "only_in_store", "only_in_reference", "content_differs", "match"
    entry_count_store: int = 0
    entry_count_reference: int = 0
    details: str = ""


@dataclass
class CompareResult:
    """Full result of a store comparison."""

    store: str
    reference_store: str
    connectors_in_store: int = 0
    connectors_in_reference: int = 0
    connectors_only_in_store: list[str] = field(default_factory=list)
    connectors_only_in_reference: list[str] = field(default_factory=list)
    connectors_matching: int = 0
    connectors_differing: int = 0
    connector_diffs: list[ConnectorDiff] = field(default_factory=list)
    index_diffs: list[IndexDiff] = field(default_factory=list)
    errors: list[str] = field(default_factory=list)

    @property
    def status(self) -> str:
        if self.errors:
            return "completed-with-errors"
        if (
            self.connectors_only_in_store
            or self.connectors_only_in_reference
            or self.connectors_differing
            or any(d.status != "match" for d in self.index_diffs)
        ):
            return "differences-found"
        return "match"

    def summary(self) -> str:
        return (
            f"[{self.status}] "
            f"Store: {self.connectors_in_store} connectors, "
            f"Reference: {self.connectors_in_reference} connectors. "
            f"Only in store: {len(self.connectors_only_in_store)}, "
            f"only in reference: {len(self.connectors_only_in_reference)}. "
            f"Matching: {self.connectors_matching}, "
            f"differing: {self.connectors_differing}. "
            f"Errors: {len(self.errors)}."
        )

    def to_dict(self) -> dict[str, Any]:
        """Serialize to a JSON-friendly dict."""
        return {
            "status": self.status,
            "store": self.store,
            "reference_store": self.reference_store,
            "connectors_in_store": self.connectors_in_store,
            "connectors_in_reference": self.connectors_in_reference,
            "connectors_only_in_store": self.connectors_only_in_store,
            "connectors_only_in_reference": self.connectors_only_in_reference,
            "connectors_matching": self.connectors_matching,
            "connectors_differing": self.connectors_differing,
            "connector_diffs": [
                {
                    "connector": d.connector,
                    "status": d.status,
                    "artifact_diffs": [
                        {
                            "file": a.file,
                            "status": a.status,
                            "details": a.details,
                        }
                        for a in d.artifact_diffs
                    ],
                }
                for d in self.connector_diffs
            ],
            "index_diffs": [
                {
                    "file": d.file,
                    "status": d.status,
                    "entry_count_store": d.entry_count_store,
                    "entry_count_reference": d.entry_count_reference,
                    "details": d.details,
                }
                for d in self.index_diffs
            ],
            "errors": self.errors,
        }


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


def _log_progress(msg: str, *args: object) -> None:
    """Log a progress message to both the logger and stderr."""
    logger.info(msg, *args)
    formatted = msg % args if args else msg
    print(formatted, file=sys.stderr, flush=True)


def _get_gcs_credentials_token() -> dict[str, str]:
    """Get GCS credentials as a token dict for gcsfs."""
    gcs_creds = os.environ.get("GCS_CREDENTIALS")
    if not gcs_creds:
        raise ValueError(
            "GCS_CREDENTIALS environment variable is required for GCS operations."
        )
    return json.loads(gcs_creds)


def _list_connectors_in_store(
    fs: gcsfs.GCSFileSystem,
    bucket: str,
    prefix: str,
) -> list[str]:
    """List connector names that have a ``latest/`` directory."""
    base = f"{bucket}/{prefix}{METADATA_FOLDER}/airbyte"
    pattern = f"{base}/*/latest/metadata.yaml"
    paths = fs.glob(pattern)

    connectors: list[str] = []
    for path in paths:
        # path: bucket/[prefix/]metadata/airbyte/<connector>/latest/metadata.yaml
        parts = path.split("/")
        # Find the connector name: 2 levels before "latest"
        try:
            latest_idx = parts.index("latest")
            connector = parts[latest_idx - 1]
            connectors.append(connector)
        except (ValueError, IndexError):
            logger.warning("Could not parse connector from path: %s", path)

    return sorted(connectors)


def _read_json_blob(
    fs: gcsfs.GCSFileSystem,
    path: str,
) -> dict[str, Any] | None:
    """Read and parse a JSON blob, returning None if it doesn't exist."""
    if not fs.exists(path):
        return None
    with fs.open(path, "r") as f:
        return json.load(f)


def _read_text_blob(
    fs: gcsfs.GCSFileSystem,
    path: str,
) -> str | None:
    """Read a text blob, returning None if it doesn't exist."""
    if not fs.exists(path):
        return None
    with fs.open(path, "r") as f:
        return f.read()


def _normalize_json(data: Any) -> str:
    """Serialize JSON with sorted keys for stable comparison."""
    return json.dumps(data, sort_keys=True, indent=2)


def _compare_json_content(
    store_data: dict[str, Any],
    ref_data: dict[str, Any],
) -> str:
    """Compare two JSON dicts and return a human-readable diff summary.

    Returns an empty string if they match.
    """
    store_normalized = _normalize_json(store_data)
    ref_normalized = _normalize_json(ref_data)

    if store_normalized == ref_normalized:
        return ""

    # Find top-level key differences
    store_keys = set(store_data.keys())
    ref_keys = set(ref_data.keys())

    diffs: list[str] = []
    only_in_store = store_keys - ref_keys
    only_in_ref = ref_keys - store_keys
    if only_in_store:
        diffs.append(f"keys only in store: {sorted(only_in_store)}")
    if only_in_ref:
        diffs.append(f"keys only in reference: {sorted(only_in_ref)}")

    # Check value differences for shared keys
    changed_keys: list[str] = []
    for key in sorted(store_keys & ref_keys):
        if _normalize_json(store_data[key]) != _normalize_json(ref_data[key]):
            changed_keys.append(key)
    if changed_keys:
        diffs.append(f"values differ for keys: {changed_keys}")

    return "; ".join(diffs) if diffs else "content differs (nested)"


# ---------------------------------------------------------------------------
# Per-connector comparison
# ---------------------------------------------------------------------------


def _compare_connector_artifacts(
    fs: gcsfs.GCSFileSystem,
    connector: str,
    store_base: str,
    ref_base: str,
) -> ConnectorDiff:
    """Compare artifacts for a single connector between store and reference."""
    artifact_diffs: list[ArtifactDiff] = []

    for filename in _ARTIFACT_FILES:
        store_path = f"{store_base}/{connector}/latest/{filename}"
        ref_path = f"{ref_base}/{connector}/latest/{filename}"

        store_exists = fs.exists(store_path)
        ref_exists = fs.exists(ref_path)

        if not store_exists and not ref_exists:
            continue
        if store_exists and not ref_exists:
            artifact_diffs.append(
                ArtifactDiff(
                    connector=connector,
                    file=filename,
                    status="only_in_store",
                )
            )
            continue
        if not store_exists and ref_exists:
            artifact_diffs.append(
                ArtifactDiff(
                    connector=connector,
                    file=filename,
                    status="only_in_reference",
                )
            )
            continue

        # Both exist -- compare contents
        if filename.endswith(".json"):
            store_data = _read_json_blob(fs, store_path)
            ref_data = _read_json_blob(fs, ref_path)
            if store_data is not None and ref_data is not None:
                details = _compare_json_content(store_data, ref_data)
                if details:
                    artifact_diffs.append(
                        ArtifactDiff(
                            connector=connector,
                            file=filename,
                            status="content_differs",
                            details=details,
                        )
                    )
        else:
            # YAML or other text -- byte comparison
            store_text = _read_text_blob(fs, store_path)
            ref_text = _read_text_blob(fs, ref_path)
            if store_text != ref_text:
                artifact_diffs.append(
                    ArtifactDiff(
                        connector=connector,
                        file=filename,
                        status="content_differs",
                        details="text content differs",
                    )
                )

    status = "match" if not artifact_diffs else "artifacts_differ"
    return ConnectorDiff(
        connector=connector,
        status=status,
        artifact_diffs=artifact_diffs,
    )


# ---------------------------------------------------------------------------
# Index comparison
# ---------------------------------------------------------------------------


def _compare_index_file(
    fs: gcsfs.GCSFileSystem,
    filename: str,
    store_bucket_root: str,
    ref_bucket_root: str,
) -> IndexDiff:
    """Compare a global registry index file between store and reference."""
    store_path = f"{store_bucket_root}/{filename}"
    ref_path = f"{ref_bucket_root}/{filename}"

    store_data = _read_json_blob(fs, store_path)
    ref_data = _read_json_blob(fs, ref_path)

    if store_data is None and ref_data is None:
        return IndexDiff(file=filename, status="match")
    if store_data is not None and ref_data is None:
        return IndexDiff(file=filename, status="only_in_store")
    if store_data is None and ref_data is not None:
        return IndexDiff(file=filename, status="only_in_reference")

    # Both exist
    assert store_data is not None
    assert ref_data is not None

    store_sources = len(store_data.get("sources", []))
    store_dests = len(store_data.get("destinations", []))
    ref_sources = len(ref_data.get("sources", []))
    ref_dests = len(ref_data.get("destinations", []))

    store_count = store_sources + store_dests
    ref_count = ref_sources + ref_dests

    details = _compare_json_content(store_data, ref_data)
    status = "match" if not details else "content_differs"

    return IndexDiff(
        file=filename,
        status=status,
        entry_count_store=store_count,
        entry_count_reference=ref_count,
        details=details,
    )


# ---------------------------------------------------------------------------
# Main compare function
# ---------------------------------------------------------------------------


def compare_stores(
    store_bucket: str,
    store_prefix: str,
    reference_bucket: str,
    reference_prefix: str,
    connector_filter: list[str] | None = None,
    *,
    with_artifacts: bool = True,
    with_indexes: bool = True,
) -> CompareResult:
    """Compare two registry stores and return structured differences.

    Args:
        store_bucket: Bucket name for the store being evaluated.
        store_prefix: Path prefix within the store bucket (e.g. ``"dev/20260306/"``).
            Must include trailing ``/`` if non-empty.
        reference_bucket: Bucket name for the known-good reference store.
        reference_prefix: Path prefix within the reference bucket.
            Must include trailing ``/`` if non-empty.
        connector_filter: If provided, only compare these connector names.
        with_artifacts: If *True* (default), compare per-connector artifact
            files (``metadata.yaml``, ``cloud.json``, ``oss.json``, ``spec.json``).
        with_indexes: If *True* (default), compare global registry index
            files (``cloud_registry.json``, ``oss_registry.json``).

    Returns:
        A :class:`CompareResult` with all differences found.
    """
    store_label = f"{store_bucket}/{store_prefix}" if store_prefix else store_bucket
    ref_label = (
        f"{reference_bucket}/{reference_prefix}"
        if reference_prefix
        else reference_bucket
    )

    result = CompareResult(store=store_label, reference_store=ref_label)

    token = _get_gcs_credentials_token()
    fs = gcsfs.GCSFileSystem(token=token)

    # Step 1: List connectors in both stores
    _log_progress("Listing connectors in store: %s", store_label)
    store_connectors = set(_list_connectors_in_store(fs, store_bucket, store_prefix))
    _log_progress("Found %d connectors in store", len(store_connectors))

    _log_progress("Listing connectors in reference: %s", ref_label)
    ref_connectors = set(
        _list_connectors_in_store(fs, reference_bucket, reference_prefix)
    )
    _log_progress("Found %d connectors in reference", len(ref_connectors))

    # Apply connector filter
    if connector_filter:
        filter_set = set(connector_filter)
        store_connectors = store_connectors & filter_set
        ref_connectors = ref_connectors & filter_set

    # Set counts based on the (possibly filtered) connector sets used for comparison
    result.connectors_in_store = len(store_connectors)
    result.connectors_in_reference = len(ref_connectors)

    # Connectors only in one store
    result.connectors_only_in_store = sorted(store_connectors - ref_connectors)
    result.connectors_only_in_reference = sorted(ref_connectors - store_connectors)

    # Add ConnectorDiff entries for connectors only in one side
    for c in result.connectors_only_in_store:
        result.connector_diffs.append(
            ConnectorDiff(connector=c, status="only_in_store")
        )
    for c in result.connectors_only_in_reference:
        result.connector_diffs.append(
            ConnectorDiff(connector=c, status="only_in_reference")
        )

    # Step 2: Compare per-connector artifacts (if requested)
    common = sorted(store_connectors & ref_connectors)

    if with_artifacts:
        _log_progress(
            "Comparing artifacts for %d connectors present in both stores...",
            len(common),
        )

        store_base = f"{store_bucket}/{store_prefix}{METADATA_FOLDER}/airbyte"
        ref_base = f"{reference_bucket}/{reference_prefix}{METADATA_FOLDER}/airbyte"

        for i, connector in enumerate(common):
            if (i + 1) % 50 == 0:
                _log_progress("  Compared %d / %d connectors...", i + 1, len(common))

            try:
                diff = _compare_connector_artifacts(fs, connector, store_base, ref_base)
            except Exception as exc:
                result.errors.append(f"Error comparing {connector}: {exc}")
                continue

            if diff.status == "match":
                result.connectors_matching += 1
            else:
                result.connectors_differing += 1
                result.connector_diffs.append(diff)

        _log_progress("Connector artifact comparison complete.")
    else:
        _log_progress("Skipping per-connector artifact comparison (--no-artifacts).")
        result.connectors_matching = len(common)

    # Step 3: Compare global index files (if requested)
    if with_indexes:
        store_root = (
            f"{store_bucket}/{store_prefix.rstrip('/')}"
            if store_prefix
            else store_bucket
        )
        ref_root = (
            f"{reference_bucket}/{reference_prefix.rstrip('/')}"
            if reference_prefix
            else reference_bucket
        )

        _log_progress("Comparing global index files...")
        for index_file in _INDEX_FILES:
            try:
                index_diff = _compare_index_file(fs, index_file, store_root, ref_root)
                result.index_diffs.append(index_diff)
            except Exception as exc:
                result.errors.append(f"Error comparing index {index_file}: {exc}")
    else:
        _log_progress("Skipping global index comparison (--no-indexes).")

    _log_progress(result.summary())
    return result
