from __future__ import annotations
import os
import time
import logging
import unicodedata
import re
import threading
from typing import Optional, Tuple, List
from enum import Enum

from fmatch.utils_normalize import strip_formerly_known_as, split_formerly_known_as

logger = logging.getLogger(__name__)

# Thread safety for file writes
_write_lock = threading.Lock()

try:
    import yaml  # type: ignore
except Exception:
    yaml = None

# Make tldextract first-class with proper PSL warming
try:
    import tldextract  # type: ignore

    # Warm the PSL cache on import
    _ = tldextract.extract("example.com")
    PSL_MODE = "psl"
    logger.info("Domain normalization mode: PSL (tldextract)")
except Exception as e:
    tldextract = None
    PSL_MODE = "fallback"
    logger.warning(
        f"Domain normalization mode: fallback (tldextract not available: {e})"
    )

PERSONAL_DOMAINS = {
    "gmail.com",
    "yahoo.com",
    "outlook.com",
    "aol.com",
    "hotmail.com",
    "icloud.com",
    "msn.com",
    "live.com",
    "yandex.com",
    "qq.com",
    "protonmail.com",
}

# Configuration constants
REGISTRY_TTL_SECONDS = 300
CONFIDENCE_EXACT_MATCH = 0.90
CONFIDENCE_NORMALIZED_MATCH = 0.85
CONFIDENCE_THRESHOLD_ROUTING = 0.90  # Minimum confidence for routing decisions

REGIONAL_TOKENS = {
    "emea",
    "apac",
    "anz",
    "na",
    "latam",
    "dach",
    "benelux",
    "eu",
    "uk",
    "us",
    "intl",
    "international",
    "global",
    "hq",
    "region",
    "western",
    "eastern",
}

CORP_SUFFIXES = {
    "inc",
    "inc.",
    "llc",
    "ltd",
    "ltd.",
    "co",
    "co.",
    "corp",
    "corp.",
    "gmbh",
    "sa",
    "sas",
    "bv",
    "plc",
    "ag",
    "oy",
    "ab",
    "srl",
    "spa",
    "pte",
}

_POSSESSIVE_FIX_RE = re.compile("\b([a-z0-9]+)(?:'|\u2019)s\b")


def domain_normalize(s: Optional[str]) -> Optional[str]:
    """
    Return registrable domain from email/URL/host or None.
    Uses tldextract if available; otherwise a conservative fallback.
    """
    if not s:
        return None
    x = str(s).strip().lower()
    if "@" in x:
        x = x.split("@", 1)[1]
    x = x.replace("https://", "").replace("http://", "").replace("www.", "")
    x = x.split("/", 1)[0].split("?", 1)[0].split("#", 1)[0]
    if tldextract:
        ext = tldextract.extract(x)
        if not ext.registered_domain:
            return None
        return ext.registered_domain
    # Fallback: keep last two labels
    parts = [p for p in x.split(".") if p]
    if len(parts) < 2:
        return None
    return ".".join(parts[-2:])


class DomainRelation(str, Enum):
    EXACT = "EXACT"
    FAMILY = "FAMILY"
    NONE = "NONE"
    # SUBDOMAIN removed - unreachable with current domain_normalize implementation


class DomainFamilyRegistry:
    """
    Lightweight registry: seed YAML + (optional) DB overlay in the future.
    Exposes alias_to_family map (registrable domains only).
    """

    def __init__(
        self,
        seed_path: Optional[str] = None,
        ttl_s: int = 300,
        learned_path: Optional[str] = None,
    ):
        self.seed_path = seed_path
        self.learned_path = learned_path or os.getenv(
            "DOMAIN_LEARNED_PATH", "config/learned_domain_mappings.yaml"
        )
        self.ttl_s = ttl_s
        self._last_load = 0.0
        self.alias_to_family: dict[str, str] = {}
        self.raw_name_to_family: dict[str, str] = {}  # exact names (lowercase)
        self.name_to_family: dict[str, str] = {}  # normalized names
        self.learned_mappings: dict[str, dict] = {}  # Track learned mappings separately
        self.display_names: dict[str, set[str]] = {}
        self.historical_names: dict[str, set[str]] = {}
        self.family_display_name: dict[str, Optional[str]] = {}
        self.family_parent_domain: dict[str, str] = {}
        self.refresh(force=True)

    def refresh(self, force: bool = False) -> bool:
        """Refresh registry from seed file. Returns True if actually refreshed."""
        if not force and (time.time() - self._last_load) < self.ttl_s:
            logger.debug("Registry refresh skipped (TTL not expired)")
            return False

        # Store count before refresh
        count_before = len(self.alias_to_family)

        self.alias_to_family.clear()
        self.raw_name_to_family.clear()
        self.name_to_family.clear()
        self.display_names.clear()
        self.historical_names.clear()
        self.family_display_name.clear()
        self.family_parent_domain.clear()
        data = []
        if self.seed_path and yaml and os.path.exists(self.seed_path):
            with open(self.seed_path, "r", encoding="utf-8") as f:
                data = yaml.safe_load(f) or []

        # Load families from seed
        families_loaded = set()
        for row in data:
            fam = str(row.get("family_id") or "").strip().lower()
            parent = domain_normalize(row.get("parent"))
            if fam and parent:
                self.alias_to_family[parent] = fam
                self.family_parent_domain[fam] = parent
                families_loaded.add(fam)
            for a in row.get("aliases") or []:
                dn = domain_normalize(a)
                if fam and dn:
                    self.alias_to_family[dn] = fam

            # Add name loading:
            names_iter = row.get("names") or []
            historical_iter = list(row.get("historical_names") or [])
            display_label = row.get("display_name")
            if display_label:
                cleaned_display = (
                    strip_formerly_known_as(display_label) or display_label.strip()
                )
                if cleaned_display:
                    self.family_display_name[fam] = cleaned_display

            for name in names_iter:
                if not name:
                    continue
                display, parsed_historical = split_formerly_known_as(name)
                candidate_names = []
                if display:
                    candidate_names.append(display)
                    self.display_names.setdefault(fam, set()).add(display)
                for hist in parsed_historical:
                    if hist:
                        self.historical_names.setdefault(fam, set()).add(hist)
                        historical_iter.append(hist)
                for candidate in candidate_names:
                    raw = candidate.strip().lower()
                    norm = self._normalize_name(candidate)
                    if raw:
                        self.raw_name_to_family[raw] = fam
                    if norm:
                        self.name_to_family[norm] = fam

            for hist in historical_iter:
                if not hist:
                    continue
                hist_clean = strip_formerly_known_as(hist) or hist.strip()
                if not hist_clean:
                    continue
                self.historical_names.setdefault(fam, set()).add(hist_clean)
                raw_hist = hist_clean.lower()
                norm_hist = self._normalize_name(hist_clean)
                if raw_hist:
                    self.raw_name_to_family[raw_hist] = fam
                if norm_hist:
                    self.name_to_family[norm_hist] = fam

            if fam not in self.family_display_name:
                if self.display_names.get(fam):
                    try:
                        self.family_display_name[fam] = sorted(
                            self.display_names[fam], key=lambda x: x.lower()
                        )[0]
                    except Exception:
                        self.family_display_name[fam] = next(
                            iter(self.display_names[fam])
                        )
                elif self.historical_names.get(fam):
                    try:
                        self.family_display_name[fam] = sorted(
                            self.historical_names[fam], key=lambda x: x.lower()
                        )[0]
                    except Exception:
                        self.family_display_name[fam] = next(
                            iter(self.historical_names[fam])
                        )

        # Load learned mappings (overlay on top of seed data)
        self._load_learned_mappings()

        # Overlay Wikidata family links (40K+ mappings, don't override manual entries)
        wikidata_enabled = (
            os.getenv("FM_COMPANY_TO_DOMAIN_WIKIDATA", "true").lower() == "true"
        )
        if wikidata_enabled:
            try:
                from ..operators import wikidata_profiles

                wikidata_links = wikidata_profiles.get_all_family_links()

                overlaid_count = 0
                self._family_to_domains = getattr(self, "_family_to_domains", {}) or {}
                for domain, family_id in wikidata_links.items():
                    if domain not in self.alias_to_family:
                        self.alias_to_family[domain] = family_id
                        self._family_to_domains.setdefault(family_id, set()).add(domain)
                        overlaid_count += 1

                logger.info(
                    "Overlaid %d Wikidata family links (total domains now: %d)",
                    overlaid_count,
                    len(self.alias_to_family),
                )
            except Exception as e:
                logger.warning("Failed to overlay Wikidata family links: %s", e)

        # Build family_to_domains mapping from alias_to_family
        family_to_domains = {}
        for domain, family in self.alias_to_family.items():
            family_to_domains.setdefault(family, set()).add(domain)
        self._family_to_domains = family_to_domains

        self._last_load = time.time()

        # Log the refresh result
        count_after = len(self.alias_to_family)
        logger.info(
            "Registry refreshed: families=%d, domains=%d->%d, seed=%s",
            len(families_loaded),
            count_before,
            count_after,
            self.seed_path or "none",
        )
        return True

    def relation(
        self, d1: Optional[str], d2: Optional[str]
    ) -> Tuple[DomainRelation, Optional[str]]:
        n1, n2 = domain_normalize(d1), domain_normalize(d2)
        if not n1 or not n2:
            return DomainRelation.NONE, None
        if n1 in PERSONAL_DOMAINS or n2 in PERSONAL_DOMAINS:
            return DomainRelation.NONE, None
        if n1 == n2:
            return DomainRelation.EXACT, self.alias_to_family.get(n1)

        # Note: SUBDOMAIN detection removed as it's unreachable with current domain_normalize
        # which returns registrable domain only (foo.bar.com -> bar.com)

        f1, f2 = self.alias_to_family.get(n1), self.alias_to_family.get(n2)
        if f1 and f1 == f2:
            return DomainRelation.FAMILY, f1
        return DomainRelation.NONE, None

    def get_domain_family(self, domain: str) -> Optional[str]:
        """Get the family ID for a domain"""
        normalized = domain_normalize(domain)
        if not normalized:
            return None
        return self.alias_to_family.get(normalized)

    def get_display_name(self, family_id: str) -> Optional[str]:
        if (
            family_id in self.family_display_name
            and self.family_display_name[family_id]
        ):
            return self.family_display_name[family_id]
        if self.display_names.get(family_id):
            try:
                return sorted(self.display_names[family_id], key=lambda x: x.lower())[0]
            except Exception:
                return next(iter(self.display_names[family_id]))
        if self.historical_names.get(family_id):
            try:
                return sorted(
                    self.historical_names[family_id], key=lambda x: x.lower()
                )[0]
            except Exception:
                return next(iter(self.historical_names[family_id]))
        return None

    def _normalize_name(self, s: Optional[str]) -> str:
        """Normalize company name: remove accents, regional tokens, corp suffixes"""
        if not s:
            return ""
        # Remove accents
        s = unicodedata.normalize("NFKD", s)
        s = "".join(ch for ch in s if not unicodedata.combining(ch))
        s = s.lower()
        s = _POSSESSIVE_FIX_RE.sub(r"\1s", s)
        # Keep alphanumeric, spaces, and &
        s = re.sub(r"[^\w\s&]", " ", s)
        # Tokenize
        tokens = [t for t in re.split(r"\s+", s) if t]
        # Remove only regional/suffix tokens (keep "systems", "labs", etc)
        tokens = [
            t for t in tokens if t not in REGIONAL_TOKENS and t not in CORP_SUFFIXES
        ]
        return " ".join(tokens)

    def lookup_family_by_name(self, name: str) -> Tuple[Optional[str], float]:
        """Look up family by company name. Returns (family_id, confidence)"""
        if not name:
            return (None, 0.0)

        # Try exact match first
        raw = name.strip().lower()
        if raw in self.raw_name_to_family:
            return (self.raw_name_to_family[raw], CONFIDENCE_EXACT_MATCH)

        # Try normalized match
        norm = self._normalize_name(name)
        if norm and norm in self.name_to_family:
            return (self.name_to_family[norm], CONFIDENCE_NORMALIZED_MATCH)

        return (None, 0.0)

    def suggest_domains(self, name: str) -> List[Tuple[str, float]]:
        """Suggest domains for a company name"""
        family_id, confidence = self.lookup_family_by_name(name)
        if not family_id:
            return []

        domains = self.family_to_domains.get(family_id, set())
        # Filter out personal domains
        domains = [d for d in domains if d not in PERSONAL_DOMAINS]
        return [(d, confidence) for d in domains]

    def add_name_override(self, name: str, family_id: str) -> None:
        """Add a name->family mapping (for learning/overrides)"""
        norm = self._normalize_name(name)
        if norm:
            self.name_to_family[norm] = family_id.lower()

    def _load_learned_mappings(self) -> None:
        """Load learned mappings from persistent storage"""
        if not self.learned_path or not os.path.exists(self.learned_path):
            return

        try:
            with open(self.learned_path, "r", encoding="utf-8") as f:
                data = yaml.safe_load(f) or {}
                self.learned_mappings = data

                # Apply learned mappings
                for mapping in data.get("mappings", []):
                    company = mapping.get("company")
                    domain = mapping.get("domain")
                    family = mapping.get("family")
                    confidence = mapping.get("confidence", 0.85)

                    if company and family:
                        display, historicals = split_formerly_known_as(company)
                        base = display or company
                        base_raw = base.strip().lower()
                        norm = self._normalize_name(base)

                        if confidence >= CONFIDENCE_EXACT_MATCH and base_raw:
                            self.raw_name_to_family[base_raw] = family
                        if norm:
                            self.name_to_family[norm] = family

                        if display:
                            self.display_names.setdefault(family, set()).add(display)
                            self.family_display_name.setdefault(family, display)
                        for hist in historicals:
                            hist_clean = hist.strip()
                            if not hist_clean:
                                continue
                            hist_raw = hist_clean.lower()
                            hist_norm = self._normalize_name(hist_clean)
                            self.historical_names.setdefault(family, set()).add(
                                hist_clean
                            )
                            if hist_raw:
                                self.raw_name_to_family[hist_raw] = family
                            if hist_norm:
                                self.name_to_family[hist_norm] = family

                        if domain:
                            normalized_domain = domain_normalize(domain)
                            if normalized_domain:
                                self.alias_to_family[normalized_domain] = family

                logger.info(
                    f"Loaded {len(data.get('mappings', []))} learned mappings from {self.learned_path}"
                )
        except Exception as e:
            logger.warning(
                f"Failed to load learned mappings from {self.learned_path}: {e}"
            )

    def _save_learned_mappings(self) -> None:
        """Persist learned mappings to storage with atomic write"""
        if not self.learned_path:
            return

        with _write_lock:
            tmp_path = f"{self.learned_path}.tmp.{os.getpid()}"

            try:
                # Ensure directory exists
                os.makedirs(os.path.dirname(self.learned_path) or ".", exist_ok=True)

                # Write to temp file
                with open(tmp_path, "w", encoding="utf-8") as f:
                    yaml.dump(self.learned_mappings, f, default_flow_style=False)

                # Atomic replace
                os.replace(tmp_path, self.learned_path)

                logger.info(
                    f"Saved {len(self.learned_mappings.get('mappings', []))} learned mappings to {self.learned_path}"
                )

            except Exception as e:
                logger.error(
                    f"Failed to save learned mappings to {self.learned_path}: {e}"
                )

            finally:
                # Clean up temp file if exists
                if os.path.exists(tmp_path):
                    try:
                        os.unlink(tmp_path)
                    except:
                        pass

    def record_verified_mapping(
        self, company_name: str, domain: str, confidence: float = 0.95
    ) -> None:
        """
        Record a verified company->domain mapping from successful matches.
        This can be used to learn from user confirmations.

        Args:
            company_name: The company name that was matched
            domain: The verified domain for this company
            confidence: Confidence level of the verification (default 0.95)
        """
        if not company_name or not domain:
            return

        normalized_domain = domain_normalize(domain)
        if not normalized_domain:
            return

        # Find or create family for this domain
        existing_family = self.alias_to_family.get(normalized_domain)
        if not existing_family:
            # Create a new family based on the domain
            existing_family = normalized_domain.replace(".", "_")
            self.alias_to_family[normalized_domain] = existing_family
            logger.info(
                f"Created new family '{existing_family}' for learned domain {normalized_domain}"
            )

        # Add the company name to this family
        display, historicals = split_formerly_known_as(company_name)
        base = display or company_name
        raw = base.strip().lower()
        norm = self._normalize_name(base)

        if raw and confidence >= CONFIDENCE_EXACT_MATCH:
            self.raw_name_to_family[raw] = existing_family
        if norm:
            self.name_to_family[norm] = existing_family

        if display:
            self.display_names.setdefault(existing_family, set()).add(display)
            self.family_display_name.setdefault(existing_family, display)
        for hist in historicals:
            hist_clean = hist.strip()
            if not hist_clean:
                continue
            hist_raw = hist_clean.lower()
            hist_norm = self._normalize_name(hist_clean)
            self.historical_names.setdefault(existing_family, set()).add(hist_clean)
            if hist_raw:
                self.raw_name_to_family[hist_raw] = existing_family
            if hist_norm:
                self.name_to_family[hist_norm] = existing_family

        # Update family_to_domains
        if hasattr(self, "_family_to_domains"):
            self._family_to_domains.setdefault(existing_family, set()).add(
                normalized_domain
            )

        # Add to learned mappings for persistence
        if "mappings" not in self.learned_mappings:
            self.learned_mappings["mappings"] = []

        # Check if already exists
        existing = False
        for m in self.learned_mappings["mappings"]:
            if (
                m.get("company") == company_name
                and m.get("domain") == normalized_domain
            ):
                m["confidence"] = max(m.get("confidence", 0), confidence)
                m["updated_at"] = time.strftime("%Y-%m-%d %H:%M:%S")
                existing = True
                break

        if not existing:
            self.learned_mappings["mappings"].append(
                {
                    "company": company_name,
                    "domain": normalized_domain,
                    "family": existing_family,
                    "confidence": confidence,
                    "created_at": time.strftime("%Y-%m-%d %H:%M:%S"),
                }
            )

        # Save to disk
        self._save_learned_mappings()

        logger.info(
            f"Learned mapping: '{company_name}' -> {normalized_domain} (family: {existing_family}, confidence: {confidence})"
        )

    def add_override(
        self, domain: str, family_id: str, account_id: Optional[str] = None
    ) -> None:
        """Add a domain override (for future tenant-specific overrides)"""
        normalized = domain_normalize(domain)
        if normalized and family_id:
            self.alias_to_family[normalized] = family_id.lower()
            logger.info(
                f"Added domain override: {normalized} -> {family_id} (account: {account_id})"
            )

    @property
    def family_to_domains(self) -> dict[str, set[str]]:
        """Get mapping from family to domains"""
        if not hasattr(self, "_family_to_domains"):
            # Build on first access if not exists
            self._family_to_domains = {}
            for domain, family in self.alias_to_family.items():
                self._family_to_domains.setdefault(family, set()).add(domain)
        return self._family_to_domains

    def reload(self, force: bool = False) -> bool:
        """Alias for refresh (backward compatibility)"""
        return self.refresh(force=force)

    def to_dict(self) -> dict:
        """Export registry as dictionary"""
        families = []
        for family_id, domains in self.family_to_domains.items():
            families.append(
                {
                    "family_id": family_id,
                    "parent": list(domains)[0]
                    if domains
                    else None,  # First domain as parent
                    "aliases": list(domains)[1:] if len(domains) > 1 else [],
                }
            )
        return {
            "families": families,
            "total_families": len(self.family_to_domains),
            "total_domains": len(self.alias_to_family),
        }


_REGISTRY: Optional[DomainFamilyRegistry] = None


def set_registry(registry: DomainFamilyRegistry) -> None:
    """Set the global registry instance."""
    global _REGISTRY
    _REGISTRY = registry


def get_registry() -> DomainFamilyRegistry:
    """Get registry instance - uses DuckDB if available, falls back to YAML"""
    global _REGISTRY

    # Check if we should use DuckDB
    use_duckdb = os.getenv("USE_DUCKDB_REGISTRY", "true").lower() == "true"
    db_path = os.getenv("DOMAIN_DB_PATH", "config/foundrygraph.duckdb")

    # Force reload if TTL expired
    if (
        _REGISTRY
        and (time.time() - getattr(_REGISTRY, "_last_load", 0)) > REGISTRY_TTL_SECONDS
    ):
        _REGISTRY = None

    if use_duckdb and os.path.exists(db_path):
        # Import and use DuckDB version
        try:
            from .domain_family_duckdb import DuckDBDomainRegistry

            if _REGISTRY is None or not isinstance(_REGISTRY, DuckDBDomainRegistry):
                _REGISTRY = DuckDBDomainRegistry(db_path)
                logger.info(f"Using DuckDB registry from {db_path}")
        except ImportError as e:
            logger.warning(f"Could not import DuckDB registry: {e}")
            # Fall back to YAML
            use_duckdb = False

    if not use_duckdb or not os.path.exists(db_path):
        # Use existing YAML implementation
        # Check by class name to avoid importing DuckDB C-extension when not needed
        is_duckdb = (
            _REGISTRY.__class__.__name__ == "DuckDBDomainRegistry" if _REGISTRY else False
        )

        if _REGISTRY is None or is_duckdb:
            seed = os.getenv("DOMAIN_FAMILY_SEED", "config/domain_families.yaml")
            _REGISTRY = DomainFamilyRegistry(seed_path=seed)
            logger.info(f"Using YAML registry from {seed}")

    _REGISTRY.refresh()
    return _REGISTRY
