"""Distributed knowledge-graph merge and cross-device entity resolution."""

from __future__ import annotations

import re
from typing import Any


class EntityResolver:
    """Resolve likely entity identity matches across local/remote graphs."""

    EXACT_MATCH = 1.0
    HIGH_CONFIDENCE = 0.85
    MEDIUM_CONFIDENCE = 0.65

    def resolve(
        self,
        local_entities: list[dict[str, Any]],
        remote_entities: list[dict[str, Any]],
    ) -> list[dict[str, Any]]:
        resolutions: list[dict[str, Any]] = []
        matched_remote: set[int] = set()

        for local in local_entities:
            best_match = None
            best_score = 0.0

            for idx, remote in enumerate(remote_entities):
                if idx in matched_remote:
                    continue
                score = self._compute_similarity(local, remote)
                if score > best_score:
                    best_score = score
                    best_match = (idx, remote)

            if best_match and best_score >= self.MEDIUM_CONFIDENCE:
                idx, remote = best_match
                matched_remote.add(idx)
                resolutions.append(
                    {
                        "local": local,
                        "remote": remote,
                        "confidence": best_score,
                        "action": "merge" if best_score >= self.HIGH_CONFIDENCE else "review",
                        "canonical_name": self._pick_canonical_name(local, remote),
                    }
                )

        for idx, remote in enumerate(remote_entities):
            if idx in matched_remote:
                continue
            resolutions.append(
                {
                    "local": None,
                    "remote": remote,
                    "confidence": 0.0,
                    "action": "add",
                    "canonical_name": str(remote.get("name", "")),
                }
            )

        return resolutions

    def _compute_similarity(self, a: dict[str, Any], b: dict[str, Any]) -> float:
        name_a = str(a.get("name", "")).lower().strip()
        name_b = str(b.get("name", "")).lower().strip()

        if str(a.get("entity_type", "")) != str(b.get("entity_type", "")):
            return 0.0
        if name_a == name_b:
            return self.EXACT_MATCH

        norm_a = self._normalize_name(name_a)
        norm_b = self._normalize_name(name_b)
        if norm_a == norm_b:
            return 0.95

        if self._is_abbreviation(name_a, name_b):
            return 0.90

        if name_a in name_b or name_b in name_a:
            return 0.80

        tokens_a = set(name_a.split())
        tokens_b = set(name_b.split())
        if tokens_a and tokens_b:
            overlap = len(tokens_a & tokens_b)
            union = len(tokens_a | tokens_b)
            jaccard = overlap / union if union > 0 else 0.0
            if jaccard >= 0.5:
                return 0.65 + jaccard * 0.2

        return 0.0

    @staticmethod
    def _normalize_name(name: str) -> str:
        name = name.lower().strip()
        name = re.sub(r"[^\w\s]", "", name)
        name = re.sub(r"\s+", " ", name)
        return name

    @staticmethod
    def _is_abbreviation(a: str, b: str) -> bool:
        parts_a = a.split()
        parts_b = b.split()
        if len(parts_a) != len(parts_b):
            return False

        for token_a, token_b in zip(parts_a, parts_b):
            if token_a == token_b:
                continue
            if (len(token_a) <= 2 and token_b.startswith(token_a[0])) or (
                len(token_b) <= 2 and token_a.startswith(token_b[0])
            ):
                continue
            return False
        return True

    @staticmethod
    def _pick_canonical_name(a: dict[str, Any], b: dict[str, Any]) -> str:
        name_a = str(a.get("name", ""))
        name_b = str(b.get("name", ""))
        return name_a if len(name_a) >= len(name_b) else name_b


class KnowledgeGraphMerger:
    """Merge remote entities/relationships into local graph state."""

    def __init__(self, storage, resolver: EntityResolver | None = None):
        self._storage = storage
        self._resolver = resolver or EntityResolver()

    def merge_entities(self, remote_entities: list[dict[str, Any]]) -> dict[str, int]:
        local_entities = self._storage.get_all_entities_as_dicts()
        resolutions = self._resolver.resolve(local_entities, remote_entities)
        stats = {"merged": 0, "added": 0, "conflicts": 0, "skipped": 0}

        for resolution in resolutions:
            action = resolution.get("action")
            if action == "add":
                self._storage.insert_remote_entity(resolution["remote"])
                stats["added"] += 1
                continue
            if action == "merge":
                self._merge_entity(
                    resolution["local"],
                    resolution["remote"],
                    str(resolution.get("canonical_name", "")),
                )
                stats["merged"] += 1
                continue
            if action == "review":
                self._storage.store_entity_conflict(resolution["local"], resolution["remote"])
                stats["conflicts"] += 1
                continue
            stats["skipped"] += 1

        return stats

    def merge_relationships(self, remote_relationships: list[dict[str, Any]], entity_id_map: dict[str, str]) -> int:
        merged = 0
        for rel in remote_relationships:
            source_id = entity_id_map.get(str(rel.get("source_entity_id", "")), str(rel.get("source_entity_id", "")))
            target_id = entity_id_map.get(str(rel.get("target_entity_id", "")), str(rel.get("target_entity_id", "")))
            self._storage.upsert_relationship_by_id(
                source_entity_id=source_id,
                target_entity_id=target_id,
                relation_type=str(rel.get("relation_type", "related-to")),
                strength=float(rel.get("strength", 0.5)),
            )
            merged += 1
        return merged

    def _merge_entity(self, local: dict[str, Any], remote: dict[str, Any], canonical_name: str) -> None:
        self._storage.update_entity(
            entity_id=local["id"],
            name=canonical_name,
            importance=max(float(local.get("importance", 0.0)), float(remote.get("importance", 0.0))),
            mention_count=int(local.get("mention_count", 0)) + int(remote.get("mention_count", 0)),
        )
