"""
DuckDB-backed Domain Family Registry with in-memory caching.
Drop-in replacement for YAML-based registry with 10x faster reloads.
"""

import duckdb
import time
import threading
import logging
import os
import re
import multiprocessing as mp
from typing import Optional, Dict, Set, Tuple, List
from contextlib import contextmanager

logger = logging.getLogger(__name__)

# Keep your existing constants
CONFIDENCE_EXACT_MATCH = 0.90
CONFIDENCE_NORMALIZED_MATCH = 0.85
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")

# Import your existing domain_normalize function
from .domain_family import domain_normalize, DomainRelation
from fmatch.utils_normalize import strip_formerly_known_as, split_formerly_known_as


class DuckDBDomainRegistry:
    """High-performance domain registry using DuckDB + in-memory cache"""

    # Thread-local storage for connections
    _thread_local = threading.local()

    def __init__(self, db_path: str = None, ttl_s: int = 300):
        self.db_path = db_path or os.getenv(
            "DOMAIN_DB_PATH", "config/foundrygraph.duckdb"
        )
        self.ttl_s = ttl_s
        self._last_load = 0.0
        self._lock = threading.RLock()

        # In-memory caches (same as your current implementation)
        self.alias_to_family: Dict[str, str] = {}
        self.raw_name_to_family: Dict[str, str] = {}
        self.name_to_family: Dict[str, str] = {}
        self.learned_mappings: Dict[str, dict] = {}
        self._family_to_domains: Dict[str, Set[str]] = {}
        self.family_display_name: Dict[str, str] = {}
        self.display_names: Dict[str, Set[str]] = {}
        self.historical_names: Dict[str, Set[str]] = {}

        # Try to initialize schema once with write access
        self._init_once()

        # Load initial cache
        self._load_cache()

    def _init_once(self):
        """One-time initialization with write access if possible"""
        try:
            if mp.current_process().name != "MainProcess":
                self.read_only = True
                return
        except Exception:
            # If multiprocessing context is unavailable, fall back to normal initialization
            pass
        try:
            # Try write access for schema creation
            conn = duckdb.connect(self.db_path, read_only=False)
            self._init_schema_with_conn(conn)
            conn.close()
            self.read_only = False
        except Exception as e:
            logger.info(f"Database opened in read-only mode: {e}")
            self.read_only = True

    @contextmanager
    def get_connection(self, read_only=True):
        """Thread-safe connection with retry logic"""
        # If a write-capable connection is already open for this thread, reuse it
        # for read operations to avoid DuckDB's single-writer locking constraints.
        if read_only and hasattr(self._thread_local, "conn_rw"):
            yield self._thread_local.conn_rw
            return

        # Use thread-local storage to avoid connection conflicts
        if read_only:
            conn_attr = "conn_ro"
            effective_read_only = True
        else:
            conn_attr = "conn_rw"
            if getattr(self, "read_only", False):
                raise RuntimeError(
                    "Write connection requested but registry is read-only"
                )
            effective_read_only = False

        if not hasattr(self._thread_local, conn_attr):
            for attempt in range(3):
                try:
                    self._thread_local.__dict__[conn_attr] = duckdb.connect(
                        self.db_path, read_only=effective_read_only
                    )
                    break
                except duckdb.IOException as e:
                    if attempt == 2:
                        # Fall back to in-memory copy for reads
                        if read_only:
                            logger.warning(f"Using in-memory database copy: {e}")
                            conn = duckdb.connect(":memory:")
                            conn.execute(f"ATTACH '{self.db_path}' AS db (READ_ONLY)")
                            self._thread_local.__dict__[conn_attr] = conn
                        else:
                            raise
                    else:
                        time.sleep(0.1 * (attempt + 1))

        yield self._thread_local.__dict__[conn_attr]

    def _init_schema_with_conn(self, conn):
        """Create tables if they don't exist using provided connection"""
        conn.execute("""
            CREATE TABLE IF NOT EXISTS domain_families (
                family_id TEXT PRIMARY KEY,
                parent_domain TEXT NOT NULL,
                display_name TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            );

            CREATE TABLE IF NOT EXISTS domain_aliases (
                domain TEXT PRIMARY KEY,
                family_id TEXT NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (family_id) REFERENCES domain_families(family_id)
            );

            CREATE TABLE IF NOT EXISTS company_names (
                raw_name TEXT NOT NULL,
                normalized_name TEXT NOT NULL,
                family_id TEXT NOT NULL,
                confidence DOUBLE DEFAULT 1.0,
                is_historical BOOLEAN DEFAULT FALSE,
                display_only BOOLEAN DEFAULT FALSE,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                PRIMARY KEY (raw_name, family_id)
            );

            CREATE SEQUENCE IF NOT EXISTS learned_mappings_seq;

            CREATE TABLE IF NOT EXISTS learned_mappings (
                id INTEGER PRIMARY KEY DEFAULT NEXTVAL('learned_mappings_seq'),
                company TEXT NOT NULL,
                domain TEXT,
                family_id TEXT NOT NULL,
                confidence DOUBLE DEFAULT 0.95,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            );

            ALTER TABLE domain_families ADD COLUMN IF NOT EXISTS display_name TEXT;
            ALTER TABLE company_names ADD COLUMN IF NOT EXISTS is_historical BOOLEAN DEFAULT FALSE;
            ALTER TABLE company_names ADD COLUMN IF NOT EXISTS display_only BOOLEAN DEFAULT FALSE;

            -- Indexes for fast lookups
            CREATE INDEX IF NOT EXISTS idx_alias_family ON domain_aliases(family_id);
            CREATE INDEX IF NOT EXISTS idx_name_normalized ON company_names(normalized_name);
            CREATE INDEX IF NOT EXISTS idx_learned_company ON learned_mappings(company);
        """)

    def _load_cache(self):
        """Load all data into memory for fast lookups"""
        start = time.perf_counter()

        with self._lock:
            with self.get_connection(read_only=True) as conn:
                # Inspect schema to see if display_name column is available
                try:
                    schema_rows = conn.execute(
                        "PRAGMA table_info('domain_families')"
                    ).fetchall()
                except Exception as schema_err:
                    logger.warning(
                        f"Could not inspect domain_families schema: {schema_err}"
                    )
                    schema_rows = []
                has_display_name = any(
                    (row[1] or "").lower() == "display_name" for row in schema_rows
                )

                # Load family metadata (tolerate missing display_name column)
                if has_display_name:
                    families = conn.execute("""
                        SELECT family_id, parent_domain, display_name
                        FROM domain_families
                    """).fetchall()
                    self.family_display_name = {
                        row[0]: (strip_formerly_known_as(row[2]) if row[2] else None)
                        for row in families
                    }
                else:
                    families = conn.execute("""
                        SELECT family_id, parent_domain
                        FROM domain_families
                    """).fetchall()
                    self.family_display_name = {row[0]: None for row in families}
                self._has_display_name_column = has_display_name

                # Load domain aliases
                aliases = conn.execute("""
                    SELECT domain, family_id FROM domain_aliases
                """).fetchall()
                self.alias_to_family = {row[0]: row[1] for row in aliases}

                try:
                    company_schema_rows = conn.execute(
                        "PRAGMA table_info('company_names')"
                    ).fetchall()
                except Exception as company_schema_err:
                    logger.warning(
                        f"Could not inspect company_names schema: {company_schema_err}"
                    )
                    company_schema_rows = []
                company_columns = {
                    (row[1] or "").lower() for row in company_schema_rows
                }
                has_is_historical = "is_historical" in company_columns
                has_display_only = "display_only" in company_columns

                # Load company names with historical flags
                name_select_parts = [
                    "raw_name",
                    "normalized_name",
                    "family_id",
                    "confidence",
                ]
                if has_is_historical:
                    name_select_parts.append("is_historical")
                if has_display_only:
                    name_select_parts.append("display_only")
                name_query = (
                    "SELECT " + ", ".join(name_select_parts) + " FROM company_names"
                )
                names = conn.execute(name_query).fetchall()

                self.raw_name_to_family.clear()
                self.name_to_family.clear()
                self.display_names.clear()
                self.historical_names.clear()

                for row in names:
                    raw, norm, family, conf = row[:4]
                    is_hist = row[4] if has_is_historical and len(row) > 4 else False
                    display_only = (
                        row[5]
                        if has_display_only and len(row) > 5
                        else row[4]
                        if (not has_is_historical and has_display_only and len(row) > 4)
                        else False
                    )

                    if isinstance(raw, str):
                        raw_lower = raw.lower()
                        sanitized = strip_formerly_known_as(raw) or raw.strip()
                    else:
                        raw_lower = raw
                        sanitized = raw

                    if (
                        not display_only
                        and conf >= CONFIDENCE_EXACT_MATCH
                        and raw_lower
                    ):
                        self.raw_name_to_family[raw_lower] = family
                        if (
                            isinstance(raw, str)
                            and sanitized
                            and sanitized.lower() != raw_lower
                        ):
                            self.raw_name_to_family[sanitized.lower()] = family
                    if norm and not display_only:
                        self.name_to_family[norm] = family
                    bucket = self.historical_names if is_hist else self.display_names
                    if isinstance(sanitized, str) and sanitized:
                        bucket.setdefault(family, set()).add(sanitized)

                # Backfill display names if table value missing
                for family, names_set in self.display_names.items():
                    if (
                        family not in self.family_display_name
                        or not self.family_display_name[family]
                    ):
                        try:
                            self.family_display_name[family] = sorted(
                                names_set, key=lambda x: x.lower()
                            )[0]
                        except Exception:
                            self.family_display_name[family] = next(iter(names_set))
                for family in self.alias_to_family.values():
                    self.family_display_name.setdefault(family, None)

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

                # Load learned mappings
                learned = conn.execute("""
                    SELECT company, domain, family_id, confidence
                    FROM learned_mappings
                    ORDER BY updated_at DESC
                """).fetchall()

                self.learned_mappings = {
                    "mappings": [
                        {
                            "company": r[0],
                            "domain": r[1],
                            "family": r[2],
                            "confidence": r[3],
                        }
                        for r in learned
                    ]
                }

                self._last_load = time.time()

        elapsed = (time.perf_counter() - start) * 1000
        logger.info(
            f"DuckDB cache loaded: {len(self.alias_to_family)} domains, "
            f"{len(self.name_to_family)} names in {elapsed:.1f}ms"
        )

    def refresh(self, force: bool = False) -> bool:
        """Refresh cache if TTL expired or forced"""
        if not force and (time.time() - self._last_load) < self.ttl_s:
            return False
        self._load_cache()
        return True

    def relation(
        self, d1: Optional[str], d2: Optional[str]
    ) -> Tuple[DomainRelation, Optional[str]]:
        """Check relationship between two domains (matches existing API)"""
        n1, n2 = domain_normalize(d1), domain_normalize(d2)
        if not n1 or not n2:
            return DomainRelation.NONE, None

        # Check for personal domains
        from .domain_family import PERSONAL_DOMAINS

        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)

        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 lookup_family_by_name(self, name: str) -> Tuple[Optional[str], float]:
        """Lookup family by company name"""
        if not name:
            return (None, 0.0)

        self.refresh()  # Check TTL

        with self._lock:
            # Try exact match
            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 lookup_family_by_domain(self, domain: str) -> Tuple[Optional[str], float]:
        """Lookup family by domain (matches existing API)"""
        norm = domain_normalize(domain)
        if not norm:
            return (None, 0.0)

        self.refresh()  # Check TTL

        with self._lock:
            family = self.alias_to_family.get(norm)
            if family:
                return (family, CONFIDENCE_EXACT_MATCH)

        return (None, 0.0)

    def lookup_families_batch(
        self, domains: list, batch_size: int = 500
    ) -> Dict[str, str]:
        """
        Batch lookup of domain families for performance.

        Args:
            domains: List of domains to lookup
            batch_size: Number of domains per SQL query (default 500)

        Returns:
            Dictionary mapping domain -> family_id
        """
        if not domains:
            return {}

        # Normalize and deduplicate domains
        normalized_map = {}  # original -> normalized
        unique_normalized = set()

        for domain in domains:
            norm = domain_normalize(domain)
            if norm:
                normalized_map[domain] = norm
                unique_normalized.add(norm)

        if not unique_normalized:
            return {}

        # Check cache first for all domains
        result = {}
        uncached = []

        with self._lock:
            for norm_domain in unique_normalized:
                if norm_domain in self.alias_to_family:
                    family = self.alias_to_family[norm_domain]
                    # Map back to all original domains with this normalized form
                    for orig, norm in normalized_map.items():
                        if norm == norm_domain:
                            result[orig] = family
                else:
                    uncached.append(norm_domain)

        # If all found in cache, return immediately
        if not uncached:
            return result

        # Batch query for uncached domains
        logger.debug(f"Batch querying {len(uncached)} uncached domains from DuckDB")

        try:
            with self.get_connection(read_only=True) as conn:
                # Process in batches to avoid SQL query size limits
                for i in range(0, len(uncached), batch_size):
                    batch = uncached[i : i + batch_size]

                    # Build parameterized query
                    placeholders = ",".join(["?" for _ in batch])
                    query = f"""
                        SELECT domain, family_id
                        FROM domain_aliases
                        WHERE domain IN ({placeholders})
                    """

                    # Execute batch query
                    rows = conn.execute(query, batch).fetchall()

                    # Process results
                    for domain, family_id in rows:
                        # Update cache
                        with self._lock:
                            self.alias_to_family[domain] = family_id

                        # Map back to original domains
                        for orig, norm in normalized_map.items():
                            if norm == domain:
                                result[orig] = family_id

                logger.debug(f"Batch lookup found {len(result)} domain families")

        except Exception as e:
            logger.warning(f"Batch lookup failed, falling back to cache only: {e}")

        return result

    def record_verified_mapping(
        self, company_name: str, domain: str, confidence: float = 0.95
    ):
        """Record a new learned mapping"""
        if not company_name or not domain:
            return

        if self.read_only:
            logger.warning("Cannot record mapping - database is read-only")
            return

        normalized_domain = domain_normalize(domain)
        if not normalized_domain:
            return

        # Find or create family
        existing_family = self.alias_to_family.get(normalized_domain)

        # Use write connection for updates
        self.cleanup_connections()
        with self.get_connection(read_only=False) as conn:
            has_display_name = getattr(self, "_has_display_name_column", False)
            if not existing_family:
                existing_family = normalized_domain.replace(".", "_")

                # Insert into DB (DuckDB compatible)
                conn.execute(
                    """
                    INSERT INTO domain_families (family_id, parent_domain)
                    SELECT ?, ?
                    WHERE NOT EXISTS (
                        SELECT 1 FROM domain_families WHERE family_id = ?
                    )
                """,
                    [existing_family, normalized_domain, existing_family],
                )

                conn.execute(
                    """
                    INSERT INTO domain_aliases (domain, family_id)
                    SELECT ?, ?
                    WHERE NOT EXISTS (
                        SELECT 1 FROM domain_aliases WHERE domain = ?
                    )
                """,
                    [normalized_domain, existing_family, normalized_domain],
                )

            # Record the mapping
            # First check if it exists
            result = conn.execute(
                """
                SELECT confidence FROM learned_mappings
                WHERE company = ? AND domain = ?
            """,
                [company_name, normalized_domain],
            ).fetchone()

            if result:
                # Update if new confidence is higher
                if confidence > result[0]:
                    conn.execute(
                        """
                        UPDATE learned_mappings
                        SET confidence = ?, updated_at = CURRENT_TIMESTAMP
                        WHERE company = ? AND domain = ?
                    """,
                        [confidence, company_name, normalized_domain],
                    )
            else:
                # Insert new mapping (explicitly use nextval for id)
                conn.execute(
                    """
                    INSERT INTO learned_mappings (id, company, domain, family_id, confidence)
                    VALUES (nextval('learned_mappings_seq'), ?, ?, ?, ?)
                """,
                    [company_name, normalized_domain, existing_family, confidence],
                )

            # Add company names (display + historical variants)
            display_label, historical_labels = split_formerly_known_as(company_name)
            base_label = (
                display_label
                or strip_formerly_known_as(company_name)
                or company_name.strip()
            )

            label_queue: List[Tuple[str, bool]] = []
            if base_label:
                label_queue.append((base_label, False))
            for hist in historical_labels:
                if hist:
                    label_queue.append((hist, True))

            if not label_queue and base_label:
                label_queue.append((base_label, False))

            if base_label and has_display_name:
                try:
                    conn.execute(
                        """
                    UPDATE domain_families
                    SET display_name = ?
                    WHERE family_id = ? AND (display_name IS NULL OR display_name = '')
                """,
                        [base_label, existing_family],
                    )
                except Exception as display_err:
                    logger.warning(
                        f"Could not update domain_families.display_name: {display_err}"
                    )

            seen_raws: Set[str] = set()
            for label, is_hist in label_queue:
                raw = label.strip().lower()
                if not raw or raw in seen_raws:
                    continue
                seen_raws.add(raw)
                norm = self._normalize_name(label)

                conn.execute(
                    """
                    DELETE FROM company_names WHERE raw_name = ? AND family_id = ?
                """,
                    [raw, existing_family],
                )

                conn.execute(
                    """
                    INSERT INTO company_names (raw_name, normalized_name, family_id, confidence, is_historical, display_only)
                    VALUES (?, ?, ?, ?, ?, FALSE)
                """,
                    [raw, norm, existing_family, confidence, is_hist],
                )

            # Explicitly commit transaction
            conn.commit()

        # Close write connection before opening read connection for cache refresh
        self.cleanup_connections()

        # Refresh cache
        self._load_cache()

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

    def to_dict(self) -> dict:
        """Return registry statistics as a dictionary"""
        return {
            "total_families": len(self._family_to_domains),
            "total_domains": len(self.alias_to_family),
            "total_company_names": len(self.name_to_family),
            "total_learned": len(self.learned_mappings),
            "implementation": "DuckDB",
        }

    @property
    def family_to_domains(self) -> Dict[str, Set[str]]:
        """Get family to domains mapping"""
        return self._family_to_domains

    def cleanup_connections(self):
        """
        Close all thread-local connections.
        Call this in long-running processes to prevent connection leaks.
        """
        try:
            if hasattr(self._thread_local, "conn_ro"):
                self._thread_local.conn_ro.close()
                delattr(self._thread_local, "conn_ro")
                logger.debug("Closed read-only connection")

            if hasattr(self._thread_local, "conn_rw"):
                self._thread_local.conn_rw.close()
                delattr(self._thread_local, "conn_rw")
                logger.debug("Closed read-write connection")
        except Exception as e:
            logger.warning(f"Error during connection cleanup: {e}")

    def __del__(self):
        """Cleanup connections when object is destroyed"""
        try:
            self.cleanup_connections()
        except:
            pass  # Suppress errors during cleanup

    def refresh_cache(self):
        """Force reload of cache from database"""
        with self._lock:
            self._load_cache()
            logger.info("Cache refreshed from database")

    def _normalize_name(self, s: Optional[str]) -> str:
        """Normalize company name (matches your existing logic)"""
        if not s:
            return ""
        import unicodedata

        # 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 and filter
        tokens = [t for t in re.split(r"\s+", s) if t]
        tokens = [
            t for t in tokens if t not in REGIONAL_TOKENS and t not in CORP_SUFFIXES
        ]

        return " ".join(tokens)

    # Compatibility methods to match existing API
    def get_domain_family(self, domain: str) -> Optional[str]:
        self.refresh()
        normalized = domain_normalize(domain)
        return self.alias_to_family.get(normalized) if normalized else None

    def get_family_domains(self, family_id: str) -> Set[str]:
        self.refresh()
        return self._family_to_domains.get(family_id, set())

    def get_display_name(self, family_id: str) -> Optional[str]:
        self.refresh()
        name = self.family_display_name.get(family_id)
        if name:
            return name
        # Fallback to first known display name or historical entry
        display_set = self.display_names.get(family_id)
        if display_set:
            try:
                return sorted(display_set, key=lambda x: x.lower())[0]
            except Exception:
                return next(iter(display_set))
        hist_set = self.historical_names.get(family_id)
        if hist_set:
            try:
                return sorted(hist_set, key=lambda x: x.lower())[0]
            except Exception:
                return next(iter(hist_set))
        return None

    def suggest_domains(self, name: str) -> List[Tuple[str, float]]:
        family_id, confidence = self.lookup_family_by_name(name)
        if not family_id:
            return []
        domains = self.get_family_domains(family_id)
        return [(d, confidence) for d in domains]

    @property
    def family_to_domains(self) -> Dict[str, Set[str]]:
        self.refresh()
        return self._family_to_domains

    def enrich_company_name(self, record: dict) -> dict:
        """Enrich record with company name based on email domain"""
        if record.get("company"):
            return record  # Already has company name

        email = record.get("email", "")
        if "@" in email:
            domain = email.split("@")[1]
            family = self.get_domain_family(domain)

            if family:
                # Try to get a company name for this family
                for name, name_family in self.name_to_family.items():
                    if name_family == family:
                        record["company"] = name.title()  # Use title case
                        break

        return record

    def get_learned_mapping(self, company_name: str) -> Optional[dict]:
        """Get learned mapping for a company name"""
        return self.learned_mappings.get(company_name)
