from collections import defaultdict
from collections.abc import Iterable, Mapping, Sequence
from dataclasses import replace
from typing import Protocol

from .adapters import (
    CERTIFICATES,
    PROXY_HOSTS,
    Capability,
    CapabilityOperation,
    DNS_RECORDS,
    Endpoint,
    EndpointPlan,
    NormalizedObject,
    ObjectReference,
    Observation,
    PlanAction,
    PlanOperation,
    RetireRequest,
)
from .dependencies import order_plan_actions


ObservationGroupKey = tuple[str, str]


class ReconciliationConflict(RuntimeError):
    """Raised when observations cannot be reconciled safely."""


class PlanningConflict(RuntimeError):
    """Raised when an endpoint cannot be planned safely."""


class PlanningAdapter(Protocol):
    endpoint: Endpoint
    capabilities: frozenset[Capability]


class ReconciliationPolicy(Protocol):
    object_type: str
    mutation_operation: PlanOperation

    def reconcile(
        self,
        logical_key: str,
        observations: Sequence[Observation],
    ) -> NormalizedObject: ...

    def equivalent(
        self,
        current: NormalizedObject,
        desired: NormalizedObject,
    ) -> bool: ...


def merge_desired_with_previous(
    fresh: Iterable[NormalizedObject],
    previous: Iterable[NormalizedObject],
    policies: Mapping[str, ReconciliationPolicy] | None = None,
) -> list[NormalizedObject]:
    selected_policies = (
        DEFAULT_POLICIES
        if policies is None
        else policies
    )

    previous_by_identity = {
        (object.object_type, object.logical_key): object
        for object in previous
    }

    merged: list[NormalizedObject] = []

    for candidate in fresh:
        identity = (
            candidate.object_type,
            candidate.logical_key,
        )
        prior = previous_by_identity.get(identity)
        policy = selected_policies.get(candidate.object_type)

        merger = (
            getattr(policy, "merge_desired", None)
            if policy is not None
            else None
        )

        if prior is not None and callable(merger):
            candidate = merger(prior, candidate)

        merged.append(candidate)

    return merged


def _observation_sort_key(
    observation: Observation,
) -> tuple[str, str, str]:
    return (
        observation.endpoint_id,
        observation.object.local_key or "",
        observation.observed_at.isoformat(),
    )


def logicalize_observation_references(
    observations: Iterable[Observation],
) -> tuple[Observation, ...]:
    rows = tuple(observations)

    local_targets: dict[
        tuple[str, str, str],
        str,
    ] = {}

    for observation in rows:
        object = observation.object

        if object.local_key is None:
            continue

        identity = (
            observation.endpoint_id,
            object.object_type,
            object.local_key,
        )

        existing = local_targets.get(identity)

        if (
            existing is not None
            and existing != object.logical_key
        ):
            raise ReconciliationConflict(
                "Conflicting endpoint-local object identity for "
                f"{observation.endpoint_id} "
                f"{object.object_type} "
                f"local key {object.local_key!r}"
            )

        local_targets[identity] = object.logical_key

    logicalized: list[Observation] = []

    for observation in rows:
        object = observation.object

        references = []

        for reference in object.references:
            if reference.logical_key is not None:
                references.append(reference)
                continue

            target = (
                observation.endpoint_id,
                reference.object_type,
                reference.local_key,
            )
            logical_key = local_targets.get(target)

            if logical_key is None:
                raise ReconciliationConflict(
                    "Unresolved endpoint-local reference from "
                    f"{observation.endpoint_id} "
                    f"{object.object_type} "
                    f"{object.logical_key!r}: "
                    f"{reference.relation} -> "
                    f"{reference.object_type} "
                    f"local key {reference.local_key!r}"
                )

            references.append(
                ObjectReference(
                    relation=reference.relation,
                    object_type=reference.object_type,
                    logical_key=logical_key,
                )
            )

        logicalized.append(
            replace(
                observation,
                object=replace(
                    object,
                    references=tuple(references),
                ),
            )
        )

    return tuple(logicalized)


def group_observations(
    observations: Iterable[Observation],
) -> dict[ObservationGroupKey, tuple[Observation, ...]]:
    grouped: dict[
        ObservationGroupKey,
        list[Observation],
    ] = defaultdict(list)

    for observation in observations:
        normalized = observation.object
        if not observation.endpoint_id:
            raise RuntimeError("Observation endpoint_id cannot be empty")
        if not normalized.object_type:
            raise RuntimeError("Observation object_type cannot be empty")
        if not normalized.logical_key:
            raise RuntimeError("Observation logical_key cannot be empty")

        key = (
            normalized.object_type,
            normalized.logical_key,
        )
        grouped[key].append(observation)

    return {
        key: tuple(sorted(rows, key=_observation_sort_key))
        for key, rows in sorted(grouped.items())
    }


class DnsUnionPolicy:
    object_type = DNS_RECORDS
    mutation_operation = PlanOperation.UPDATE

    def reconcile(
        self,
        logical_key: str,
        observations: Sequence[Observation],
    ) -> NormalizedObject:
        if not observations:
            raise RuntimeError(
                "DNS union policy requires at least one observation"
            )

        first = observations[0].object
        conflicting_endpoints = {
            observation.endpoint_id
            for observation in observations
            if observation.object.payload != first.payload
        }
        if conflicting_endpoints:
            conflicting_endpoints.add(observations[0].endpoint_id)
            endpoint_list = ", ".join(
                sorted(conflicting_endpoints)
            )
            raise ReconciliationConflict(
                f"Conflicting {DNS_RECORDS} observations for "
                f"{logical_key!r} from {endpoint_list}"
            )

        return NormalizedObject(
            object_type=DNS_RECORDS,
            logical_key=logical_key,
            payload=dict(first.payload),
        )

    def equivalent(
        self,
        current: NormalizedObject,
        desired: NormalizedObject,
    ) -> bool:
        return (
            current.object_type == desired.object_type
            and current.logical_key == desired.logical_key
            and current.payload == desired.payload
            and current.references == desired.references
        )


class CertificateNewestPolicy:
    object_type = CERTIFICATES
    mutation_operation = PlanOperation.REPLACE

    _PORTABLE_FIELDS = (
        "domain_names",
        "valid_from",
        "expires_on",
    )

    @classmethod
    def _portable_state(
        cls,
        object: NormalizedObject,
    ) -> tuple[object, ...]:
        return tuple(
            object.payload.get(field)
            for field in cls._PORTABLE_FIELDS
        )

    @staticmethod
    def _date_key(
        value: object,
        field: str,
    ) -> tuple[int, str]:
        if value is None:
            return (0, "")

        if not isinstance(value, str):
            raise ReconciliationConflict(
                f"Certificate {field} must be normalized text"
            )

        return (1, value)

    @classmethod
    def _newest_key(
        cls,
        object: NormalizedObject,
    ) -> tuple[
        tuple[int, str],
        tuple[int, str],
    ]:
        return (
            cls._date_key(
                object.payload.get("valid_from"),
                "valid_from",
            ),
            cls._date_key(
                object.payload.get("expires_on"),
                "expires_on",
            ),
        )

    def reconcile(
        self,
        logical_key: str,
        observations: Sequence[Observation],
    ) -> NormalizedObject:
        if not observations:
            raise RuntimeError(
                "Certificate policy requires at least one observation"
            )

        newest_key = max(
            self._newest_key(observation.object)
            for observation in observations
        )

        newest = [
            observation
            for observation in observations
            if self._newest_key(observation.object) == newest_key
        ]

        baseline = self._portable_state(newest[0].object)

        conflicting_endpoints = {
            observation.endpoint_id
            for observation in newest
            if self._portable_state(observation.object) != baseline
        }

        if conflicting_endpoints:
            conflicting_endpoints.add(newest[0].endpoint_id)

            raise ReconciliationConflict(
                f"Conflicting {CERTIFICATES} observations for "
                f"{logical_key!r} at newest issuance from "
                + ", ".join(sorted(conflicting_endpoints))
            )

        winner = newest[0].object

        return NormalizedObject(
            object_type=CERTIFICATES,
            logical_key=logical_key,
            payload=dict(winner.payload),
        )

    def equivalent(
        self,
        current: NormalizedObject,
        desired: NormalizedObject,
    ) -> bool:
        return (
            current.object_type == CERTIFICATES
            and desired.object_type == CERTIFICATES
            and current.logical_key == desired.logical_key
            and self._portable_state(current)
            == self._portable_state(desired)
        )

    def merge_desired(
        self,
        previous: NormalizedObject,
        fresh: NormalizedObject,
    ) -> NormalizedObject:
        if (
            previous.object_type != CERTIFICATES
            or fresh.object_type != CERTIFICATES
            or previous.logical_key != fresh.logical_key
        ):
            raise ReconciliationConflict(
                "Certificate desired-state merge requires "
                "matching certificate identities"
            )

        if self._newest_key(previous) > self._newest_key(fresh):
            return previous

        return fresh

    @staticmethod
    def _equivalent_selection_key(
        object: NormalizedObject,
    ) -> tuple[int, str]:
        # Preserve a native ACME-renewable certificate over an
        # imported replica when both contain the same portable
        # certificate state.
        provider = object.payload.get("provider")

        return (
            1 if provider == "letsencrypt" else 0,
            object.local_key or "",
        )

    def select_current(
        self,
        observations: Sequence[Observation],
        desired: NormalizedObject,
    ) -> tuple[
        NormalizedObject,
        tuple[RetireRequest, ...],
    ]:
        if not observations:
            raise PlanningConflict(
                "Certificate current-state selection requires "
                "at least one observation"
            )

        candidates = [
            observation.object
            for observation in observations
        ]

        if any(
            candidate.local_key is None
            for candidate in candidates
        ):
            raise PlanningConflict(
                "Observed certificate requires endpoint-local identity"
            )

        matching = [
            candidate
            for candidate in candidates
            if self.equivalent(
                candidate,
                desired,
            )
        ]

        if matching:
            selected = max(
                matching,
                key=self._equivalent_selection_key,
            )
        else:
            # No local copy is already desired. Replace the newest
            # local certificate and retire every other stale copy
            # after the replacement and its dependents verify.
            selected = max(
                candidates,
                key=lambda candidate: (
                    self._newest_key(candidate),
                    candidate.local_key or "",
                ),
            )

        retirements = tuple(
            RetireRequest(object=candidate)
            for candidate in candidates
            if candidate.local_key != selected.local_key
        )

        return selected, retirements


class ProxyHostConsensusPolicy:
    object_type = PROXY_HOSTS
    mutation_operation = PlanOperation.UPDATE

    def reconcile(
        self,
        logical_key: str,
        observations: Sequence[Observation],
    ) -> NormalizedObject:
        if not observations:
            raise RuntimeError(
                "Proxy-host policy requires at least one observation"
            )

        first = observations[0].object

        conflicting_endpoints = {
            observation.endpoint_id
            for observation in observations
            if (
                observation.object.payload != first.payload
                or observation.object.references != first.references
            )
        }

        if conflicting_endpoints:
            conflicting_endpoints.add(observations[0].endpoint_id)

            raise ReconciliationConflict(
                f"Conflicting {PROXY_HOSTS} observations for "
                f"{logical_key!r} from "
                + ", ".join(sorted(conflicting_endpoints))
            )

        for reference in first.references:
            if reference.local_key is not None:
                raise ReconciliationConflict(
                    "Proxy-host desired state cannot contain "
                    "endpoint-local references"
                )

        return NormalizedObject(
            object_type=PROXY_HOSTS,
            logical_key=logical_key,
            payload=dict(first.payload),
            references=first.references,
        )

    def equivalent(
        self,
        current: NormalizedObject,
        desired: NormalizedObject,
    ) -> bool:
        return (
            current.object_type == PROXY_HOSTS
            and desired.object_type == PROXY_HOSTS
            and current.logical_key == desired.logical_key
            and current.payload == desired.payload
            and current.references == desired.references
        )


DEFAULT_POLICIES: Mapping[str, ReconciliationPolicy] = {
    DNS_RECORDS: DnsUnionPolicy(),
    CERTIFICATES: CertificateNewestPolicy(),
    PROXY_HOSTS: ProxyHostConsensusPolicy(),
}


def reconcile_observations(
    observations: Iterable[Observation],
    policies: Mapping[str, ReconciliationPolicy] | None = None,
) -> list[NormalizedObject]:
    selected_policies = (
        DEFAULT_POLICIES
        if policies is None
        else policies
    )
    desired: list[NormalizedObject] = []

    logicalized = logicalize_observation_references(
        observations
    )

    for (object_type, logical_key), rows in (
        group_observations(logicalized).items()
    ):
        policy = selected_policies.get(object_type)
        if policy is None:
            raise RuntimeError(
                f"No reconciliation policy for {object_type!r}"
            )
        if policy.object_type != object_type:
            raise RuntimeError(
                f"Policy for {object_type!r} advertises "
                f"{policy.object_type!r}"
            )

        desired.append(policy.reconcile(logical_key, rows))

    return desired


def build_endpoint_plans(
    adapters: Iterable[PlanningAdapter],
    desired_objects: Iterable[NormalizedObject],
    observations: Iterable[Observation],
    policies: Mapping[str, ReconciliationPolicy] | None = None,
) -> tuple[EndpointPlan, ...]:
    desired_by_identity: dict[
        tuple[str, str],
        NormalizedObject,
    ] = {}

    for desired in desired_objects:
        if desired.local_key is not None:
            raise RuntimeError(
                "Desired objects in plans cannot contain "
                "an endpoint-local key"
            )

        identity = (
            desired.object_type,
            desired.logical_key,
        )
        if identity in desired_by_identity:
            raise RuntimeError(
                "Duplicate desired object identity: "
                f"{desired.object_type} {desired.logical_key}"
            )
        desired_by_identity[identity] = desired

    selected_policies = (
        DEFAULT_POLICIES
        if policies is None
        else policies
    )

    raw_observations = tuple(observations)

    current_by_identity: dict[
        tuple[str, str, str],
        list[Observation],
    ] = defaultdict(list)

    raw_current_by_identity: dict[
        tuple[str, str, str],
        list[Observation],
    ] = defaultdict(list)

    for observation in raw_observations:
        raw_current_by_identity[
            (
                observation.endpoint_id,
                observation.object.object_type,
                observation.object.logical_key,
            )
        ].append(observation)

    logicalized_observations = (
        logicalize_observation_references(
            raw_observations
        )
    )

    for observation in logicalized_observations:
        current_by_identity[
            (
                observation.endpoint_id,
                observation.object.object_type,
                observation.object.logical_key,
            )
        ].append(observation)

    adapters_by_endpoint: dict[str, PlanningAdapter] = {}
    for adapter in adapters:
        endpoint_id = adapter.endpoint.identity
        if endpoint_id in adapters_by_endpoint:
            raise RuntimeError(
                f"Duplicate adapter endpoint: {endpoint_id}"
            )
        adapters_by_endpoint[endpoint_id] = adapter

    plans: list[EndpointPlan] = []
    for endpoint_id, adapter in sorted(
        adapters_by_endpoint.items()
    ):
        apply_types = {
            capability.object_type
            for capability in adapter.capabilities
            if capability.operation == CapabilityOperation.APPLY
        }
        if not apply_types:
            continue

        actions: list[PlanAction] = []
        planned_retirements: list[
            RetireRequest
        ] = []

        for identity, desired in sorted(
            desired_by_identity.items()
        ):
            object_type, logical_key = identity
            if object_type not in apply_types:
                continue

            current_rows = current_by_identity.get(
                (endpoint_id, object_type, logical_key),
                [],
            )
            policy = selected_policies.get(object_type)

            if len(current_rows) > 1:
                selector = (
                    getattr(
                        policy,
                        "select_current",
                        None,
                    )
                    if policy is not None
                    else None
                )

                if selector is None:
                    raise PlanningConflict(
                        f"Multiple endpoint objects for {endpoint_id} "
                        f"{object_type} {logical_key!r}"
                    )

                current, retirements = selector(
                    tuple(current_rows),
                    desired,
                )

                planned_retirements.extend(
                    retirements
                )
            else:
                current = (
                    current_rows[0].object
                    if current_rows
                    else None
                )
            if current is None:
                operation = PlanOperation.CREATE
            else:
                if policy is None:
                    equivalent = (
                        current.payload == desired.payload
                        and current.references == desired.references
                    )
                else:
                    equivalent = policy.equivalent(
                        current,
                        desired,
                    )

                operation = (
                    PlanOperation.NOOP
                    if equivalent
                    else (
                        policy.mutation_operation
                        if policy is not None
                        else PlanOperation.UPDATE
                    )
                )

            actions.append(
                PlanAction(
                    endpoint_id=endpoint_id,
                    operation=operation,
                    desired=desired,
                    current=current,
                )
            )

        replaced = {
            (
                action.desired.object_type,
                action.desired.logical_key,
            )
            for action in actions
            if (
                action.operation == PlanOperation.REPLACE
                and action.desired is not None
            )
        }

        retiring_local = {
            (
                retirement.object.object_type,
                retirement.object.local_key,
            )
            for retirement in planned_retirements
        }

        if replaced or retiring_local:
            promoted: list[PlanAction] = []

            for action in actions:
                desired = action.desired

                logical_dependency_replaced = (
                    desired is not None
                    and any(
                        (
                            reference.object_type,
                            reference.logical_key,
                        )
                        in replaced
                        for reference in desired.references
                        if reference.logical_key is not None
                    )
                )

                raw_rows = (
                    raw_current_by_identity.get(
                        (
                            endpoint_id,
                            desired.object_type,
                            desired.logical_key,
                        ),
                        [],
                    )
                    if desired is not None
                    else []
                )

                local_dependency_retiring = any(
                    (
                        reference.object_type,
                        reference.local_key,
                    )
                    in retiring_local
                    for row in raw_rows
                    for reference in row.object.references
                    if reference.local_key is not None
                )

                if (
                    action.operation == PlanOperation.NOOP
                    and (
                        logical_dependency_replaced
                        or local_dependency_retiring
                    )
                ):
                    reason = (
                        "referenced object is being replaced"
                        if logical_dependency_replaced
                        else (
                            "referenced local object is being retired"
                        )
                    )

                    action = replace(
                        action,
                        operation=PlanOperation.UPDATE,
                        reason=reason,
                    )

                promoted.append(action)

            actions = promoted

        plans.append(
            EndpointPlan(
                endpoint_id=endpoint_id,
                actions=order_plan_actions(actions),
                retirements=tuple(
                    planned_retirements
                ),
            )
        )

    return tuple(plans)