from typing import Optional, List, Dict
from datetime import datetime
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select
import logging

from ..models import (
    Investigation,
    InvestigationStatus,
    Finding,
    FindingCategory,
    Integration,
    IntegrationProvider,
)
from .salesforce_gateway import SalesforceGateway
import pandas as pd  # Still needed for DataFrame operations

# NEW: Use engine for all matching
from .engine_client import get_engine_client
from ...preprocessing.rules.domain_utils import registrable_domain

log = logging.getLogger(__name__)


class InvestigationService:
    """Service for running Salesforce investigations"""

    def __init__(self, db: AsyncSession):
        self.db = db

    async def run_investigation(self, investigation_id: str):
        """Run a Salesforce investigation"""
        # Get investigation
        investigation = await self.db.get(Investigation, investigation_id)
        if not investigation:
            raise ValueError(f"Investigation {investigation_id} not found")

        # Update status
        investigation.status = InvestigationStatus.RUNNING
        investigation.started_at = datetime.utcnow()
        await self.db.commit()

        try:
            # Get Salesforce integration
            integration = await self.db.execute(
                select(Integration).where(
                    Integration.account_id == investigation.account_id,
                    Integration.provider == IntegrationProvider.SALESFORCE,
                )
            )
            sf_integration = integration.scalar_one_or_none()

            if not sf_integration:
                raise ValueError("Salesforce not connected")

            # Create gateway for tenant
            client = await SalesforceGateway.for_tenant(
                sf_integration.account_id, self.db
            )

            # Define callback to persist token updates
            async def _save_integration(integ):
                await self.db.merge(integ)
                await self.db.commit()
                await self.db.refresh(integ)

            # Calibrate rate limits based on org limits
            # Gateway handles rate limits internally; optional calibration via client is not needed

            # Run based on params
            params = investigation.params
            objects = params.get("objects", ["Account", "Lead"])

            findings = []
            stats = {"total_records": 0, "duplicates": 0, "link_candidates": 0}

            for obj_type in objects:
                if obj_type == "Account":
                    await self._scan_accounts(client, investigation_id, findings, stats)
                elif obj_type == "Lead":
                    await self._scan_leads(client, investigation_id, findings, stats)
                elif obj_type == "Contact":
                    await self._scan_contacts(client, investigation_id, findings, stats)

            # Save findings
            for finding_data in findings:
                finding = Finding(**finding_data)
                self.db.add(finding)

            # Update investigation
            investigation.status = InvestigationStatus.COMPLETED
            investigation.finished_at = datetime.utcnow()
            investigation.stats = stats
            await self.db.commit()

            log.info(f"Investigation {investigation_id} completed: {stats}")

        except Exception as e:
            log.error(f"Investigation {investigation_id} failed: {e}")
            investigation.status = InvestigationStatus.FAILED
            investigation.finished_at = datetime.utcnow()
            await self.db.commit()
            raise

    async def _scan_accounts(
        self, client, investigation_id: str, findings: List[Dict], stats: Dict
    ):
        """Scan for duplicate accounts using the engine"""

        # Define callback to persist token updates
        async def _save_integration(integ):
            await self.db.merge(integ)
            await self.db.commit()
            await self.db.refresh(integ)

        # Query accounts
        query = """
            SELECT Id, Name, Website, Phone, BillingCity, BillingState
            FROM Account
            WHERE IsDeleted = false
            LIMIT 1000
        """

        result = await client.soql(query)
        records = result.get("records", [])
        stats["total_records"] += len(records)
        if not records:
            return

        df = pd.DataFrame(records)
        # Normalized columns for engine
        df["WebsiteDomain"] = df["Website"].map(self._domain_from_url)
        df["PhoneDigits"] = df["Phone"].map(self._digits)
        df["NameClean"] = df["Name"].fillna("").str.lower()

        eng = get_engine_client()

        mappings = [
            {
                "source_column": "NameClean",
                "reference_column": "NameClean",
                "weight": 0.7,
                "algorithm": "Auto",
            },
            {
                "source_column": "WebsiteDomain",
                "reference_column": "WebsiteDomain",
                "weight": 0.3,
                "algorithm": "Domain",
            },
            # Optional: {"source_column": "PhoneDigits", "reference_column": "PhoneDigits", "weight": 0.2, "algorithm": "Exact"},
        ]

        res = await eng.find_duplicates(
            data=df[["Id", "NameClean", "WebsiteDomain", "PhoneDigits"]],
            mappings=mappings,
            # TODO: Replace with intelligence API
            # threshold=0.80  # HARDCODED - NEEDS FIX,                # same spirit as before
            id_column="Id",
            blocking_strategy={"domain": True, "name_trigram": True},
            job_id=None,
        )

        # Convert pairwise matches -> clusters (union-find)
        pairs = (
            res.matches[["id1", "id2", "match_score"]].values.tolist()
            if not res.matches.empty
            else []
        )
        clusters = self._cluster_pairs(pairs)

        for idx, member_ids in enumerate(clusters, start=1):
            cluster_records = df[df["Id"].isin(member_ids)][
                ["Id", "Name", "Website"]
            ].to_dict("records")
            findings.append(
                {
                    "investigation_id": investigation_id,
                    "object_type": "Account",
                    "category": FindingCategory.DUPLICATE,
                    "score": 0.9,  # or use avg of pair scores in cluster
                    "cluster_id": f"cluster_{idx}",
                    "records": cluster_records,
                    "suggestion": {"primary": member_ids[0]},
                }
            )
            stats["duplicates"] += 1

    async def _scan_leads(
        self, client, investigation_id: str, findings: List[Dict], stats: Dict
    ):
        """Scan for leads that should be linked to accounts using the engine"""

        # Define callback to persist token updates
        async def _save_integration(integ):
            await self.db.merge(integ)
            await self.db.commit()
            await self.db.refresh(integ)

        # Get leads
        lead_query = """
            SELECT Id, Email, Company, Phone, Name, Website
            FROM Lead
            WHERE IsConverted = false AND IsDeleted = false
            LIMIT 1000
        """

        leads = await client.soql(lead_query)
        lead_records = leads.get("records", [])

        # Get accounts using bulk API for larger datasets
        account_query = """
            SELECT Id, Name, Website, Phone
            FROM Account
            WHERE IsDeleted = false
        """

        # Use bulk API if expecting many records
        if len(lead_records) > 100:
            account_records = await client.bulk_query(account_query, "Account")
        else:
            accounts = await client.soql(account_query)
            account_records = accounts.get("records", [])

        if not lead_records or not account_records:
            return

        leads_df = pd.DataFrame(lead_records)
        accts_df = pd.DataFrame(account_records)

        # Prep normalized columns
        leads_df["EmailDomain"] = leads_df["Email"].map(self._domain_from_email)
        leads_df["CompanyClean"] = leads_df["Company"].fillna("").str.lower()
        leads_df["PhoneDigits"] = (
            leads_df["Phone"].map(self._digits) if "Phone" in leads_df else ""
        )

        accts_df["WebsiteDomain"] = accts_df["Website"].map(self._domain_from_url)
        accts_df["NameClean"] = accts_df["Name"].fillna("").str.lower()
        accts_df["PhoneDigits"] = (
            accts_df["Phone"].map(self._digits) if "Phone" in accts_df else ""
        )

        eng = get_engine_client()
        mappings = [
            {
                "source_column": "EmailDomain",
                "reference_column": "WebsiteDomain",
                "weight": 0.60,
                "algorithm": "Domain",
            },
            {
                "source_column": "CompanyClean",
                "reference_column": "NameClean",
                "weight": 0.30,
                "algorithm": "Auto",
            },
            # {"source_column":"PhoneDigits", "reference_column":"PhoneDigits", "weight":0.10, "algorithm":"Exact"},
        ]

        res = await eng.match_records(
            source_data=leads_df[["Id", "EmailDomain", "CompanyClean"]],
            reference_data=accts_df[["Id", "WebsiteDomain", "NameClean"]],
            mappings=mappings,
            # TODO: Replace with intelligence API
            # threshold=0.75  # HARDCODED - NEEDS FIX,
            source_id_column="Id",
            reference_id_column="Id",
            blocking_strategy={"domain": True, "name_trigram": True},
        )

        if res.matches.empty:
            return

        # Expect columns: source_id, target_id, match_score, field_scores, explanation
        for _, row in res.matches.iterrows():
            lead_id = row["source_id"]
            acct_id = row["target_id"]
            lead_rec = (
                leads_df[leads_df["Id"] == lead_id].iloc[0]
                if not leads_df[leads_df["Id"] == lead_id].empty
                else {}
            )
            acct_rec = (
                accts_df[accts_df["Id"] == acct_id].iloc[0]
                if not accts_df[accts_df["Id"] == acct_id].empty
                else {}
            )

            findings.append(
                {
                    "investigation_id": investigation_id,
                    "object_type": "Lead",
                    "category": FindingCategory.LINK_CANDIDATE,
                    "score": float(row.get("match_score", 0.0)),
                    "records": [
                        {
                            "Id": lead_id,
                            "Email": lead_rec.get("Email", ""),
                            "Company": lead_rec.get("Company", ""),
                        },
                        {
                            "Id": acct_id,
                            "Name": acct_rec.get("Name", ""),
                            "Website": acct_rec.get("Website", ""),
                        },
                    ],
                    "suggestion": {"lead_id": lead_id, "account_id": acct_id},
                }
            )
            stats["link_candidates"] += 1

    async def _scan_contacts(
        self,
        client: SalesforceGateway,
        investigation_id: str,
        findings: List[Dict],
        stats: Dict,
    ):
        """Scan for duplicate contacts using the engine"""

        # Define callback to persist token updates
        async def _save_integration(integ):
            await self.db.merge(integ)
            await self.db.commit()
            await self.db.refresh(integ)

        query = """
            SELECT Id, FirstName, LastName, Email, Phone, AccountId,
                   Account.Name, Account.Website
            FROM Contact
            WHERE IsDeleted = false
            LIMIT 5000
        """

        contacts = []
        # Use paginate_soql if available, otherwise use regular soql
        try:
            async for batch in client.paginate_soql(query, _save_integration):
                contacts.extend(batch)
        except AttributeError:
            # Fall back to regular soql if paginate_soql doesn't exist
            result = await client.soql(query, _save_integration)
            contacts = result.get("records", [])

        stats["total_records"] += len(contacts)
        if not contacts:
            return

        df = pd.DataFrame(contacts)
        df["EmailLower"] = df["Email"].fillna("").str.lower()
        df["NameFull"] = (
            (df["FirstName"].fillna("") + " " + df["LastName"].fillna(""))
            .str.strip()
            .str.lower()
        )

        eng = get_engine_client()
        mappings = [
            {
                "source_column": "EmailLower",
                "reference_column": "EmailLower",
                "weight": 3.0,
                "algorithm": "Exact",
            },
            {
                "source_column": "NameFull",
                "reference_column": "NameFull",
                "weight": 1.5,
                "algorithm": "TokenSetRatio",
            },
            # Same-account bonus is better applied in engine as a custom feature, but
            # if not available, you can post-filter on AccountId equality when forming clusters.
        ]

        res = await eng.find_duplicates(
            data=df[["Id", "EmailLower", "NameFull", "AccountId"]],
            mappings=mappings,
            # TODO: Replace with intelligence API
            # threshold=0.85  # HARDCODED - NEEDS FIX,
            id_column="Id",
            blocking_strategy={"email": True, "name_trigram": True},
        )

        pairs = (
            res.matches[["id1", "id2", "match_score"]].values.tolist()
            if not res.matches.empty
            else []
        )
        clusters = self._cluster_pairs(pairs)

        for idx, member_ids in enumerate(clusters, start=1):
            cluster_records = df[df["Id"].isin(member_ids)][
                ["Id", "FirstName", "LastName", "Email", "AccountId"]
            ].to_dict("records")
            findings.append(
                {
                    "investigation_id": investigation_id,
                    "object_type": "Contact",
                    "category": FindingCategory.DUPLICATE,
                    "score": 0.9,
                    "cluster_id": f"contact_cluster_{idx}",
                    "records": cluster_records,
                    "suggestion": {"primary": member_ids[0]},
                }
            )
            stats["duplicates"] += 1

        # Find contacts without Account linkage
        orphaned = df[df["AccountId"].isna()]

        for _, contact in orphaned.iterrows():
            if not contact.get("Email"):
                continue

            email_domain = self._extract_domain(contact["Email"])
            if self._is_personal_domain(email_domain):
                continue

            # Try to find matching Account
            account_query = f"""
                SELECT Id, Name, Website
                FROM Account
                WHERE Website LIKE '%{email_domain}%'
                LIMIT 1
            """

            try:
                accounts = await client.soql(account_query, _save_integration)
                if accounts.get("records"):
                    account = accounts["records"][0]

                    # Strip unnecessary attributes to keep findings light
                    clean_contact = {
                        "Id": contact["Id"],
                        "FirstName": contact.get("FirstName"),
                        "LastName": contact.get("LastName"),
                        "Email": contact.get("Email"),
                    }
                    clean_account = {
                        "Id": account["Id"],
                        "Name": account.get("Name"),
                        "Website": account.get("Website"),
                    }

                    findings.append(
                        {
                            "investigation_id": investigation_id,
                            "object_type": "Contact",
                            "category": FindingCategory.LINK_CANDIDATE,
                            "score": 0.85,
                            "records": [clean_contact, clean_account],
                            "suggestion": {
                                "contact_id": contact["Id"],
                                "account_id": account["Id"],
                                "field": "AccountId",
                            },
                        }
                    )

                    stats["link_candidates"] += 1
            except Exception as e:
                log.warning(f"Failed to find Account for Contact {contact['Id']}: {e}")
                pass

    def _extract_domain(self, url_or_email: str) -> Optional[str]:
        """Extract domain from URL or email with better normalization"""
        if not url_or_email:
            return None

        # Handle email
        if "@" in url_or_email:
            domain = url_or_email.split("@")[-1].lower()
        else:
            # Handle URL
            from urllib.parse import urlparse

            try:
                parsed = urlparse(
                    url_or_email
                    if url_or_email.startswith("http")
                    else f"http://{url_or_email}"
                )
                domain = parsed.netloc.lower()
            except:
                return None

        # Strip common prefixes
        for prefix in ("www.", "links.", "click.", "email.", "go."):
            if domain.startswith(prefix):
                domain = domain[len(prefix) :]

        # Handle IDN/punycode domains
        try:
            import idna

            domain = idna.decode(domain)
        except:
            pass  # Keep as-is if not punycode

        return domain

    def _compare_domains(self, url1: str, url2: str) -> float:
        """Compare two URLs/domains"""
        domain1 = self._extract_domain(url1)
        domain2 = self._extract_domain(url2)

        if not domain1 or not domain2:
            return 0.0

        return 1.0 if domain1 == domain2 else 0.0

    def _is_personal_domain(self, domain: str) -> bool:
        """Check if domain is personal email"""
        personal_domains = {
            "gmail.com",
            "yahoo.com",
            "hotmail.com",
            "outlook.com",
            "aol.com",
            "icloud.com",
            "me.com",
            "mac.com",
            "live.com",
            "msn.com",
            "protonmail.com",
            "ymail.com",
        }
        return domain in personal_domains if domain else False

    def _normalize_name(self, s: str) -> str:
        """Normalize company name for bucketing"""
        s = (s or "").lower()
        # Remove common suffixes
        for suffix in (
            " inc",
            " inc.",
            " llc",
            " ltd",
            " co",
            " corp",
            " corporation",
            " company",
        ):
            s = s.replace(suffix, "")
        return " ".join(s.split())

    def _bucket_accounts(self, df: pd.DataFrame) -> pd.core.groupby.DataFrameGroupBy:
        """Bucket accounts by domain or name prefix for efficient duplicate detection"""

        def bucket_key(row):
            # Try domain first
            dom = self._extract_domain(row.get("Website", ""))
            if dom and not self._is_personal_domain(dom):
                return f"dom:{dom}"
            # Fall back to first 6 chars of normalized name
            name = self._normalize_name(row.get("Name", ""))
            return f"nm:{name[:6]}" if name else "nm:unknown"

        df["__bucket"] = df.apply(bucket_key, axis=1)
        return df.groupby("__bucket")

    # NEW HELPER FUNCTIONS FOR ENGINE-BASED MATCHING
    def _digits(self, s):
        """Extract only digits from a string."""
        return "".join(c for c in str(s or "") if c.isdigit())

    def _domain_from_email(self, s: str) -> str:
        """Extract domain from email address."""
        s = (s or "").strip().lower()
        return s.split("@")[-1] if "@" in s else ""

    def _domain_from_url(self, s: str) -> str:
        """Extract domain from URL."""
        # Use global domain utility for consistent eTLD+1 extraction
        return registrable_domain(s or "")

    def _cluster_pairs(self, pairs):
        """Convert pairwise matches to clusters using union-find."""
        # pairs: List[[id1, id2, score]]
        parent = {}

        def find(x):
            parent.setdefault(x, x)
            if parent[x] != x:
                parent[x] = find(parent[x])
            return parent[x]

        def union(a, b):
            ra, rb = find(a), find(b)
            if ra != rb:
                parent[rb] = ra

        for a, b, _ in pairs:
            union(a, b)

        groups = {}
        for a, b, _ in pairs:
            groups.setdefault(find(a), set()).update([a, b])
        return [sorted(list(v)) for v in groups.values()]
