from __future__ import annotations

import hashlib
import logging
import os
import re
from typing import Any, Dict, Iterable, List, Optional, Tuple

from .fg_entity_id import generate_event_id, generate_fg_id

logger = logging.getLogger(__name__)

try:
    import tldextract

    _tld_extractor = tldextract.TLDExtract(cache_dir=False, suffix_list_urls=())
except Exception:  # pragma: no cover - best effort optional dependency
    _tld_extractor = None

_DOMAIN_TWO_PART_TLDS = {
    "co.uk",
    "com.au",
    "co.jp",
    "co.nz",
    "co.in",
    "co.za",
    "com.br",
    "com.cn",
    "com.sg",
    "com.hk",
    "com.mx",
    "org.uk",
    "gov.uk",
    "ac.uk",
    "net.au",
    "org.au",
    "co.id",
    "co.kr",
    "co.th",
    "com.tw",
    "com.ar",
}

_bq_client = None


def _project_id() -> str:
    return (
        os.getenv("FM_FOUNDRYGRAPH_PROJECT_ID")
        or os.getenv("GOOGLE_CLOUD_PROJECT")
        or ""
    ).strip()


def _table_ref(dataset: str, table: str) -> str:
    project_id = _project_id()
    if project_id:
        return f"{project_id}.{dataset}.{table}"
    return f"{dataset}.{table}"


def _get_bq_client():
    global _bq_client
    project_id = _project_id()
    if _bq_client is None or (project_id and not getattr(_bq_client, "project", None)):
        from google.cloud import bigquery

        _bq_client = bigquery.Client(project=project_id) if project_id else bigquery.Client()
    return _bq_client


def _normalize_domain(value: str) -> str:
    if not value:
        return ""
    cleaned = value.strip().lower()
    cleaned = re.sub(r"^https?://", "", cleaned)
    cleaned = cleaned.split("@")[-1]
    cleaned = cleaned.split("/")[0]
    cleaned = cleaned.split("?")[0]
    cleaned = cleaned.split("#")[0]
    cleaned = cleaned.split(":")[0]
    if cleaned.startswith("www."):
        cleaned = cleaned[4:]
    cleaned = cleaned.strip(".")
    if not cleaned:
        return ""
    return _registrable_domain(cleaned)


def _registrable_domain(domain: str) -> str:
    if not domain:
        return ""
    if _tld_extractor is not None:
        try:
            extracted = _tld_extractor(domain)
        except Exception:
            extracted = None
        if extracted:
            if extracted.domain and extracted.suffix:
                return f"{extracted.domain}.{extracted.suffix}"
            if extracted.domain:
                return extracted.domain
    parts = domain.split(".")
    if len(parts) <= 2:
        return domain
    last_two = ".".join(parts[-2:])
    if last_two in _DOMAIN_TWO_PART_TLDS and len(parts) >= 3:
        return ".".join(parts[-3:])
    return last_two


class AliasNormalizer:
    """Per-system normalization for alias values."""

    @staticmethod
    def normalize(id_system: str, id_value: str) -> str:
        """Normalize alias value based on system type."""
        system = (id_system or "").lower().strip()
        value = "" if id_value is None else str(id_value).strip()
        if not value:
            return ""
        if system == "domain":
            return _normalize_domain(value)
        if system == "wikidata":
            match = re.search(r"Q\d+", value.upper())
            return match.group(0) if match else value.upper()
        if system == "lei":
            return value.upper().replace(" ", "")
        if system == "ticker":
            return value.upper().strip()
        if system in ("duns", "sec_cik", "isin"):
            return value.strip()
        return value.strip()

    @staticmethod
    def compute_alias_key(id_system: str, id_value: str) -> str:
        """Generate deterministic key for uniqueness: SHA256(system:normalized_value)."""
        system = (id_system or "").lower().strip()
        normalized = AliasNormalizer.normalize(system, id_value)
        return hashlib.sha256(f"{system}:{normalized}".encode()).hexdigest()


class InMemoryEntityStore:
    """Simple in-memory store for testing registry behavior."""

    def __init__(self) -> None:
        self.entities: Dict[str, Dict[str, Any]] = {}
        self.aliases: Dict[str, Dict[str, Any]] = {}
        self.events: List[Dict[str, Any]] = []


class EntityRegistry:
    """Central registry for fg_id resolution and management."""

    def __init__(
        self,
        *,
        store: Optional[InMemoryEntityStore] = None,
        dataset: str = "foundrygraph_gold",
    ) -> None:
        self._store = store
        self._dataset = dataset
        self._aliases_table = _table_ref(dataset, "company_aliases")
        self._companies_table = _table_ref(dataset, "companies_gold")
        self._events_table = _table_ref(dataset, "company_events")

    async def resolve_by_alias(self, id_system: str, id_value: str) -> Optional[str]:
        """Resolve external identifier to fg_id. Returns None if no mapping."""
        alias_key = AliasNormalizer.compute_alias_key(id_system, id_value)
        if self._store is not None:
            alias = self._store.aliases.get(alias_key)
            return alias.get("fg_id") if alias else None

        try:
            client = _get_bq_client()
        except Exception as exc:
            logger.warning("FG registry BQ client unavailable: %s", exc)
            return None

        try:
            from google.cloud import bigquery

            query = f"""
                SELECT fg_id
                FROM `{self._aliases_table}`
                WHERE alias_key = @alias_key
                LIMIT 1
            """
            job_config = bigquery.QueryJobConfig(
                query_parameters=[
                    bigquery.ScalarQueryParameter("alias_key", "STRING", alias_key)
                ]
            )
            rows = list(client.query(query, job_config=job_config, project=_project_id() or None).result())
            if rows:
                return getattr(rows[0], "fg_id", None)
        except Exception as exc:
            logger.warning("FG registry resolve_by_alias failed: %s", exc)
        return None

    async def bulk_resolve(
        self,
        identifiers: List[Tuple[str, str]],
    ) -> Dict[Tuple[str, str], Optional[str]]:
        """Batch resolve with a single BigQuery query when possible."""
        alias_key_map: Dict[Tuple[str, str], str] = {}
        for system, value in identifiers or []:
            alias_key_map[(system, value)] = AliasNormalizer.compute_alias_key(system, value)
        if self._store is not None:
            out: Dict[Tuple[str, str], Optional[str]] = {}
            for key, alias_key in alias_key_map.items():
                alias = self._store.aliases.get(alias_key)
                out[key] = alias.get("fg_id") if alias else None
            return out

        if not alias_key_map:
            return {}

        try:
            client = _get_bq_client()
        except Exception as exc:
            logger.warning("FG registry BQ client unavailable: %s", exc)
            return {k: None for k in alias_key_map}

        try:
            from google.cloud import bigquery

            alias_keys = sorted(set(alias_key_map.values()))
            query = f"""
                SELECT alias_key, fg_id
                FROM `{self._aliases_table}`
                WHERE alias_key IN UNNEST(@alias_keys)
            """
            job_config = bigquery.QueryJobConfig(
                query_parameters=[
                    bigquery.ArrayQueryParameter("alias_keys", "STRING", alias_keys)
                ]
            )
            rows = list(client.query(query, job_config=job_config, project=_project_id() or None).result())
            lookup = {getattr(row, "alias_key", None): getattr(row, "fg_id", None) for row in rows}
        except Exception as exc:
            logger.warning("FG registry bulk_resolve failed: %s", exc)
            lookup = {}

        resolved: Dict[Tuple[str, str], Optional[str]] = {}
        for key, alias_key in alias_key_map.items():
            resolved[key] = lookup.get(alias_key)
        return resolved

    async def resolve_or_mint(
        self,
        aliases: List[Tuple[str, str]],
        entity_data: Optional[Dict[str, Any]] = None,
        job_id: Optional[str] = None,
    ) -> Tuple[str, bool]:
        """
        Resolve existing fg_id or mint new one.
        Returns (fg_id, was_new).
        Normalizes all aliases at write-time.
        """
        normalized_aliases = []
        for system, value in aliases or []:
            system_norm = (system or "").lower().strip()
            if not system_norm or value is None:
                continue
            normalized_value = AliasNormalizer.normalize(system_norm, value)
            if not normalized_value:
                continue
            normalized_aliases.append((system_norm, normalized_value))
        if not normalized_aliases:
            raise ValueError("At least one alias is required")

        if self._store is not None:
            existing_ids = {
                self._store.aliases.get(AliasNormalizer.compute_alias_key(system, value), {}).get("fg_id")
                for system, value in normalized_aliases
            }
            existing_ids.discard(None)
            if existing_ids:
                if len(existing_ids) > 1:
                    raise ValueError("Alias conflict: multiple fg_ids for provided aliases")
                return existing_ids.pop(), False

            fg_id = generate_fg_id("co")
            self._store.entities[fg_id] = {
                "fg_id": fg_id,
                "status": "active",
                "redirect_to_fg_id": None,
                "entity_type": "company",
            }
            primary_set = False
            for system, value in normalized_aliases:
                alias_key = AliasNormalizer.compute_alias_key(system, value)
                is_primary = system == "wikidata" or not primary_set
                primary_set = primary_set or is_primary
                self._store.aliases[alias_key] = {
                    "fg_id": fg_id,
                    "id_system": system,
                    "id_value": value,
                    "alias_key": alias_key,
                    "is_primary": is_primary,
                    "source": "registry",
                    "job_id": job_id,
                    "entity_type": "company",
                }
            self._store.events.append(
                {
                    "event_id": generate_event_id(),
                    "subject_fg_id": fg_id,
                    "event_type": "created",
                    "job_id": job_id,
                }
            )
            return fg_id, True

        # BigQuery-backed flow
        existing_map = await self.bulk_resolve([(s, v) for s, v in normalized_aliases])
        existing_ids = {fg for fg in existing_map.values() if fg}
        if existing_ids:
            if len(existing_ids) > 1:
                raise ValueError("Alias conflict: multiple fg_ids for provided aliases")
            return existing_ids.pop(), False

        fg_id = generate_fg_id("co")
        wikidata_id = None
        if entity_data:
            wikidata_id = entity_data.get("wikidata_id")

        if wikidata_id:
            await self._update_company_fg_id(wikidata_id, fg_id)

        await self._merge_aliases(fg_id, normalized_aliases, job_id=job_id)
        await self._insert_event(
            {
                "event_id": generate_event_id(),
                "subject_fg_id": fg_id,
                "event_type": "created",
                "job_id": job_id,
                "reason": None,
                "actor": "system",
            }
        )

        return fg_id, True

    async def get_entity(
        self,
        fg_id: str,
        follow_redirects: bool = True,
    ) -> Optional[Dict[str, Any]]:
        """
        Get entity record by fg_id.

        If follow_redirects=True and status='deprecated':
            - Follow redirect_to_fg_id to survivor
            - Return survivor profile + deprecated_fg_id field

        If follow_redirects=False:
            - Return record as-is with status and redirect_to fields
        """
        if self._store is not None:
            record = self._store.entities.get(fg_id)
            if not record:
                return None
            record = dict(record)
            if follow_redirects and record.get("status") == "deprecated" and record.get("redirect_to_fg_id"):
                survivor = self._store.entities.get(record.get("redirect_to_fg_id"))
                if not survivor:
                    return record
                survivor_copy = dict(survivor)
                survivor_copy["deprecated_fg_id"] = fg_id
                return survivor_copy
            return record

        try:
            client = _get_bq_client()
        except Exception as exc:
            logger.warning("FG registry BQ client unavailable: %s", exc)
            return None

        try:
            from google.cloud import bigquery

            query = f"""
                SELECT fg_id, wikidata_id, label, domain, status, redirect_to_fg_id
                FROM `{self._companies_table}`
                WHERE fg_id = @fg_id
                LIMIT 1
            """
            job_config = bigquery.QueryJobConfig(
                query_parameters=[
                    bigquery.ScalarQueryParameter("fg_id", "STRING", fg_id)
                ]
            )
            rows = list(client.query(query, job_config=job_config, project=_project_id() or None).result())
            if not rows:
                return None
            row = rows[0]
            record = {
                "fg_id": getattr(row, "fg_id", None),
                "wikidata_id": getattr(row, "wikidata_id", None),
                "label": getattr(row, "label", None),
                "domain": getattr(row, "domain", None),
                "status": getattr(row, "status", None) or "active",
                "redirect_to_fg_id": getattr(row, "redirect_to_fg_id", None),
            }
        except Exception as exc:
            logger.warning("FG registry get_entity failed: %s", exc)
            return None

        if follow_redirects and record.get("status") == "deprecated" and record.get("redirect_to_fg_id"):
            survivor = await self.get_entity(record["redirect_to_fg_id"], follow_redirects=False)
            if survivor:
                survivor = dict(survivor)
                survivor["deprecated_fg_id"] = fg_id
                return survivor
        return record

    async def merge_entities(
        self,
        survivor_fg_id: str,
        victim_fg_id: str,
        reason: str,
        actor: str,
        job_id: Optional[str] = None,
    ) -> bool:
        """
        Merge victim into survivor.

        Steps:
        1. Validate both exist and are active
        2. Update victim: status='deprecated', redirect_to_fg_id=survivor
        3. Migrate aliases from victim to survivor (update fg_id)
        4. Log MERGE event with merged_from_fg_id + merged_into_fg_id
        """
        if self._store is not None:
            survivor = self._store.entities.get(survivor_fg_id)
            victim = self._store.entities.get(victim_fg_id)
            if not survivor or not victim:
                return False
            if survivor.get("status") != "active" or victim.get("status") != "active":
                return False
            victim["status"] = "deprecated"
            victim["redirect_to_fg_id"] = survivor_fg_id
            for alias in self._store.aliases.values():
                if alias.get("fg_id") == victim_fg_id:
                    alias["fg_id"] = survivor_fg_id
            self._store.events.append(
                {
                    "event_id": generate_event_id(),
                    "subject_fg_id": survivor_fg_id,
                    "event_type": "deprecated",
                    "merged_from_fg_id": victim_fg_id,
                    "merged_into_fg_id": survivor_fg_id,
                    "reason": reason,
                    "actor": actor,
                    "job_id": job_id,
                }
            )
            return True

        # BigQuery-backed merge
        try:
            client = _get_bq_client()
        except Exception as exc:
            logger.warning("FG registry BQ client unavailable: %s", exc)
            return False

        try:
            from google.cloud import bigquery

            status_query = f"""
                SELECT fg_id, status
                FROM `{self._companies_table}`
                WHERE fg_id IN UNNEST(@ids)
            """
            status_config = bigquery.QueryJobConfig(
                query_parameters=[
                    bigquery.ArrayQueryParameter("ids", "STRING", [survivor_fg_id, victim_fg_id])
                ]
            )
            status_rows = list(client.query(status_query, job_config=status_config, project=_project_id() or None).result())
            status_map = {getattr(r, "fg_id", None): getattr(r, "status", None) for r in status_rows}
            if survivor_fg_id not in status_map or victim_fg_id not in status_map:
                return False
            if (status_map.get(survivor_fg_id) or "active") != "active":
                return False
            if (status_map.get(victim_fg_id) or "active") != "active":
                return False

            update_victim = f"""
                UPDATE `{self._companies_table}`
                SET status = 'deprecated',
                    redirect_to_fg_id = @survivor,
                    fg_id_updated_at = CURRENT_TIMESTAMP()
                WHERE fg_id = @victim
            """
            update_config = bigquery.QueryJobConfig(
                query_parameters=[
                    bigquery.ScalarQueryParameter("survivor", "STRING", survivor_fg_id),
                    bigquery.ScalarQueryParameter("victim", "STRING", victim_fg_id),
                ]
            )
            client.query(update_victim, job_config=update_config, project=_project_id() or None).result()

            alias_update = f"""
                UPDATE `{self._aliases_table}`
                SET fg_id = @survivor, updated_at = CURRENT_TIMESTAMP()
                WHERE fg_id = @victim
            """
            alias_config = bigquery.QueryJobConfig(
                query_parameters=[
                    bigquery.ScalarQueryParameter("survivor", "STRING", survivor_fg_id),
                    bigquery.ScalarQueryParameter("victim", "STRING", victim_fg_id),
                ]
            )
            client.query(alias_update, job_config=alias_config, project=_project_id() or None).result()

            await self._insert_event(
                {
                    "event_id": generate_event_id(),
                    "subject_fg_id": survivor_fg_id,
                    "event_type": "deprecated",
                    "merged_from_fg_id": victim_fg_id,
                    "merged_into_fg_id": survivor_fg_id,
                    "reason": reason,
                    "actor": actor,
                    "job_id": job_id,
                }
            )
            return True
        except Exception as exc:
            logger.warning("FG registry merge_entities failed: %s", exc)
            return False

    async def _update_company_fg_id(self, wikidata_id: str, fg_id: str) -> None:
        try:
            client = _get_bq_client()
        except Exception as exc:
            logger.warning("FG registry BQ client unavailable: %s", exc)
            return

        try:
            from google.cloud import bigquery

            query = f"""
                UPDATE `{self._companies_table}`
                SET fg_id = @fg_id,
                    fg_id_created_at = CURRENT_TIMESTAMP(),
                    fg_id_updated_at = CURRENT_TIMESTAMP()
                WHERE wikidata_id = @wikidata_id
                  AND fg_id IS NULL
            """
            job_config = bigquery.QueryJobConfig(
                query_parameters=[
                    bigquery.ScalarQueryParameter("fg_id", "STRING", fg_id),
                    bigquery.ScalarQueryParameter("wikidata_id", "STRING", wikidata_id),
                ]
            )
            client.query(query, job_config=job_config, project=_project_id() or None).result()
        except Exception as exc:
            logger.warning("FG registry update_company_fg_id failed: %s", exc)

    async def _merge_aliases(
        self,
        fg_id: str,
        aliases: Iterable[Tuple[str, str]],
        job_id: Optional[str] = None,
    ) -> None:
        try:
            client = _get_bq_client()
        except Exception as exc:
            logger.warning("FG registry BQ client unavailable: %s", exc)
            return

        primary_set = False
        for system, value in aliases:
            alias_key = AliasNormalizer.compute_alias_key(system, value)
            is_primary = system == "wikidata" or not primary_set
            primary_set = primary_set or is_primary
            try:
                from google.cloud import bigquery

                query = f"""
                    MERGE `{self._aliases_table}` T
                    USING (
                        SELECT
                            @fg_id AS fg_id,
                            @id_system AS id_system,
                            @id_value AS id_value,
                            @alias_key AS alias_key,
                            @is_primary AS is_primary,
                            @source AS source,
                            @job_id AS job_id,
                            @entity_type AS entity_type
                    ) S
                    ON T.alias_key = S.alias_key
                    WHEN NOT MATCHED THEN
                      INSERT (fg_id, id_system, id_value, alias_key, is_primary, source, job_id, entity_type)
                      VALUES (S.fg_id, S.id_system, S.id_value, S.alias_key, S.is_primary, S.source, S.job_id, S.entity_type)
                """
                job_config = bigquery.QueryJobConfig(
                    query_parameters=[
                        bigquery.ScalarQueryParameter("fg_id", "STRING", fg_id),
                        bigquery.ScalarQueryParameter("id_system", "STRING", system),
                        bigquery.ScalarQueryParameter("id_value", "STRING", value),
                        bigquery.ScalarQueryParameter("alias_key", "STRING", alias_key),
                        bigquery.ScalarQueryParameter("is_primary", "BOOL", is_primary),
                        bigquery.ScalarQueryParameter("source", "STRING", "registry"),
                        bigquery.ScalarQueryParameter("job_id", "STRING", job_id),
                        bigquery.ScalarQueryParameter("entity_type", "STRING", "company"),
                    ]
                )
                client.query(query, job_config=job_config, project=_project_id() or None).result()
            except Exception as exc:
                logger.warning("FG registry alias merge failed: %s", exc)

    async def _insert_event(self, event: Dict[str, Any]) -> None:
        try:
            client = _get_bq_client()
        except Exception as exc:
            logger.warning("FG registry BQ client unavailable: %s", exc)
            return

        try:
            from google.cloud import bigquery

            query = f"""
                INSERT `{self._events_table}`
                (event_id, subject_fg_id, event_type, merged_from_fg_id, merged_into_fg_id, reason, actor, job_id)
                VALUES (@event_id, @subject_fg_id, @event_type, @merged_from_fg_id, @merged_into_fg_id, @reason, @actor, @job_id)
            """
            job_config = bigquery.QueryJobConfig(
                query_parameters=[
                    bigquery.ScalarQueryParameter("event_id", "STRING", event.get("event_id")),
                    bigquery.ScalarQueryParameter("subject_fg_id", "STRING", event.get("subject_fg_id")),
                    bigquery.ScalarQueryParameter("event_type", "STRING", event.get("event_type")),
                    bigquery.ScalarQueryParameter("merged_from_fg_id", "STRING", event.get("merged_from_fg_id")),
                    bigquery.ScalarQueryParameter("merged_into_fg_id", "STRING", event.get("merged_into_fg_id")),
                    bigquery.ScalarQueryParameter("reason", "STRING", event.get("reason")),
                    bigquery.ScalarQueryParameter("actor", "STRING", event.get("actor")),
                    bigquery.ScalarQueryParameter("job_id", "STRING", event.get("job_id")),
                ]
            )
            client.query(query, job_config=job_config, project=_project_id() or None).result()
        except Exception as exc:
            logger.warning("FG registry insert event failed: %s", exc)
