from __future__ import annotations

from collections import Counter, defaultdict
from contextlib import asynccontextmanager
import hashlib
import json
import logging
import re
import time
import uuid
from datetime import datetime, timezone
from typing import Dict, Iterable, List, Optional, Tuple

from sqlalchemy.ext.asyncio import AsyncSession

from fmatch.saas.repos.salesforce_schema_index_repo import SalesforceSchemaIndexRepo
from fmatch.saas.model_defs.salesforce_schema_index import (
    SalesforceSchemaField,
    SalesforceSchemaObject,
    SalesforceSchemaRelationship,
    SalesforceSchemaTerm,
)
from fmatch.saas.services.salesforce_gateway import SalesforceGateway
from fmatch.saas import settings as saas_settings
from fmatch.saas.services.telemetry import emit

log = logging.getLogger(__name__)


DEFAULT_OBJECT_ALLOWLIST = [
    "Opportunity",
    "Account",
    "Contact",
    "Lead",
    "User",
    "Campaign",
    "CampaignMember",
    "OpportunityContactRole",
    "OpportunityLineItem",
    "Product2",
    "PricebookEntry",
    "Case",
    "Task",
    "Event",
]

DOMAIN_SYNONYMS = {
    "arr": ["ARR__c", "AnnualRecurringRevenue__c"],
    "carr": ["CARR__c"],
    "acv": ["ACV__c", "AnnualContractValue__c"],
    "tcv": ["TCV__c", "TotalContractValue__c"],
    "renewal": ["RenewalAmount__c"],
    "pipeline": ["Pipeline__c", "Amount"],
    "forecast": ["ForecastCategoryName", "ForecastCategory"],
}


def _coerce_utc(value: Optional[datetime]) -> Optional[datetime]:
    if value is None:
        return None
    if value.tzinfo is None:
        return value.replace(tzinfo=timezone.utc)
    return value.astimezone(timezone.utc)


class SalesforceSchemaIndexer:
    """Builds and maintains a per-org Salesforce schema index."""

    def __init__(self, db: AsyncSession, settings_module=None) -> None:
        self.repo = SalesforceSchemaIndexRepo(db)
        self.settings = settings_module or saas_settings
        self.db = db

    @asynccontextmanager
    async def _acquire_org_lock(self, org_id: str, stage: str):
        token = str(uuid.uuid4())
        try:
            await self.repo.mark_in_progress(org_id, stage=stage, token=token)
        except RuntimeError:
            raise RuntimeError("Index already running")
        try:
            yield token
        finally:
            await self.repo.clear_in_progress(org_id, token=token)

    def _resolve_fls_hash(self, gateway: SalesforceGateway) -> Optional[str]:
        candidate = getattr(gateway, "fls_scope_hash", None)
        if callable(candidate):
            try:
                candidate = candidate()
            except Exception:  # noqa: BLE001
                candidate = None
        if not candidate:
            candidate = getattr(gateway, "user_profile_id", None)
        if isinstance(candidate, str):
            return candidate
        return None

    @staticmethod
    def _normalize_object_candidates(
        global_describe: Dict[str, any],
    ) -> List[Dict[str, any]]:
        sobjects = global_describe.get("sobjects", []) or []
        # focus on queryable objects
        filtered = [
            obj
            for obj in sobjects
            if obj.get("queryable") and not obj.get("deprecatedAndHidden", False)
        ]
        return filtered

    @staticmethod
    def _calculate_schema_hash(objects: Iterable[Dict[str, any]]) -> str:
        payload = json.dumps(objects, sort_keys=True, ensure_ascii=False)
        return hashlib.sha256(payload.encode("utf-8")).hexdigest()

    async def ensure_quick_index(
        self,
        *,
        org_id: str,
        gateway: SalesforceGateway,
        max_custom_objects: int = 5,
    ) -> str:
        """Ensure we have at least the quick index built for this org; return schema hash."""
        ttl_seconds = max(1, self.settings.AI_SCHEMA_INDEX_TTL_HOURS) * 3600
        state = await self.repo.get_state(org_id)
        state_quick_indexed_at = _coerce_utc(state.quick_indexed_at) if state else None
        if (
            state
            and state_quick_indexed_at
            and (datetime.now(timezone.utc) - state_quick_indexed_at).total_seconds()
            < ttl_seconds
        ):
            emit(
                "index.quick.ms",
                org_id=org_id,
                ms=0,
                cached=True,
                schema_hash=state.schema_hash or "",
            )
            return state.schema_hash or ""

        try:
            async with self._acquire_org_lock(org_id, stage="quick") as token:
                log.info("Building Salesforce quick schema index for org %s", org_id)
                start = time.perf_counter()
                await self.repo.touch_heartbeat(org_id, token=token)

                global_desc = await gateway.describe_global()
                candidates = self._normalize_object_candidates(global_desc)
                object_selection: List[Dict[str, any]] = []
                standard_seen = set()
                for item in candidates:
                    api_name = item.get("name")
                    if (
                        api_name in DEFAULT_OBJECT_ALLOWLIST
                        and api_name not in standard_seen
                    ):
                        object_selection.append(item)
                        standard_seen.add(api_name)

                custom_candidates = [
                    obj
                    for obj in candidates
                    if obj.get("custom") and obj.get("name") not in standard_seen
                ]
                custom_candidates.sort(
                    key=lambda o: o.get("lastModifiedDate") or "", reverse=True
                )
                object_selection.extend(custom_candidates[:max_custom_objects])

                schema_hash = self._calculate_schema_hash(
                    [
                        {
                            "name": obj.get("name"),
                            "custom": obj.get("custom", False),
                            "label": obj.get("label"),
                        }
                        for obj in object_selection
                    ]
                )

                state = await self.repo.get_state(org_id)
                fls_hash = self._resolve_fls_hash(gateway)
                fls_matches = (
                    not fls_hash
                    or not state
                    or not getattr(state, "fls_hash", None)
                    or state.fls_hash == fls_hash
                )
                if (
                    state
                    and state.schema_hash == schema_hash
                    and state.quick_indexed_at
                    and fls_matches
                ):
                    await self.repo.mark_quick_indexed(
                        org_id, schema_hash=schema_hash, fls_hash=fls_hash
                    )
                    emit(
                        "index.quick.ms",
                        org_id=org_id,
                        ms=0,
                        cached=True,
                        schema_hash=schema_hash,
                    )
                    return schema_hash
                if state and state.schema_hash == schema_hash and state.quick_indexed_at and not fls_matches:
                    log.info(
                        "FLS scope changed for org %s (old=%s new=%s); rebuilding schema index",
                        org_id,
                        state.fls_hash,
                        fls_hash,
                    )

                fields_payload: List[SalesforceSchemaField] = []
                relationships_payload: List[SalesforceSchemaRelationship] = []
                objects_payload: List[SalesforceSchemaObject] = []
                terms_payload: List[SalesforceSchemaTerm] = []

                for obj in object_selection:
                    api_name = obj.get("name")
                    try:
                        describe = await gateway.describe_object(api_name)
                    except Exception as exc:  # noqa: BLE001
                        log.warning(
                            "Failed to describe %s for org %s: %s",
                            api_name,
                            org_id,
                            exc,
                        )
                        continue

                    objects_payload.append(
                        SalesforceSchemaObject(
                            org_id=org_id,
                            object_api=api_name,
                            label=describe.get("label"),
                            label_plural=describe.get("labelPlural"),
                            is_custom=describe.get("custom", False),
                            is_reportable=describe.get("queryable", False),
                            schema_hash=schema_hash,
                            record_count=describe.get("recordCount"),
                            last_modified=None,
                        )
                    )

                    for field in describe.get("fields", []) or []:
                        is_accessible = field.get(
                            "accessible",
                            field.get("permissions", {}).get("readable", True),
                        )
                        if is_accessible is False:
                            continue
                        if not field.get("queryable", True):
                            continue
                        field_api = field.get("name")
                        data_type = field.get("type", "string")
                        picklist_values = (
                            [
                                entry.get("value")
                                for entry in field.get("picklistValues", [])
                                if entry.get("active", True)
                            ]
                            if data_type == "picklist"
                            else []
                        )
                        relationships = field.get("referenceTo") or []
                        relationships = [ref for ref in relationships if ref]

                        fields_payload.append(
                            SalesforceSchemaField(
                                org_id=org_id,
                                object_api=api_name,
                                field_api=field_api,
                                label=field.get("label"),
                                help_text=field.get("inlineHelpText"),
                                data_type=data_type,
                                precision=field.get("precision"),
                                scale=field.get("scale"),
                                length=field.get("length"),
                                is_custom=field.get("custom", False),
                                is_formula=field.get("calculated", False),
                                is_calculated=field.get("calculated", False),
                                is_picklist=data_type == "picklist",
                                picklist_json=json.dumps(
                                    picklist_values, ensure_ascii=False
                                ),
                                reference_to=",".join(relationships)
                                if relationships
                                else None,
                                relationship_api=field.get("relationshipName"),
                                relationship_label=field.get("relationshipLabel"),
                                queryable=field.get("queryable", True),
                                filterable=field.get("filterable", True),
                                sortable=field.get("sortable", True),
                                nillable=field.get("nillable", True),
                                unique=field.get("unique", False),
                                calculated_formula=field.get("calculatedFormula"),
                                schema_hash=schema_hash,
                            )
                        )

                        base_terms = set()
                        for token_source in (
                            field.get("label"),
                            field_api,
                            field.get("inlineHelpText"),
                        ):
                            if not token_source:
                                continue
                            for token in (
                                str(token_source)
                                .replace("_", " ")
                                .replace(".", " ")
                                .split()
                            ):
                                cleaned = token.strip().lower()
                                if cleaned:
                                    base_terms.add(cleaned)
                        for alias_term, mapped_fields in DOMAIN_SYNONYMS.items():
                            if (
                                field_api in mapped_fields
                                or field.get("label") in mapped_fields
                            ):
                                base_terms.add(alias_term)

                        for term in base_terms:
                            weight = 1.0
                            if term in DOMAIN_SYNONYMS:
                                weight = 2.0
                            terms_payload.append(
                                SalesforceSchemaTerm(
                                    org_id=org_id,
                                    term=term,
                                    field_ref=f"{api_name}.{field_api}",
                                    weight=weight,
                                    source="label",
                                    schema_hash=schema_hash,
                                )
                            )

                        # Process relationships for this field
                        # Note: For polymorphic lookups (fields that reference multiple objects),
                        # we only store the first target since the UNIQUE constraint on
                        # (org_id, from_object, field_api) allows only one relationship per field.
                        # The full list of targets is stored in the fields table's reference_to column.
                        if relationships:
                            # Only store the first target (primary relationship)
                            relationships_payload.append(
                                SalesforceSchemaRelationship(
                                    org_id=org_id,
                                    from_object=api_name,
                                    to_object=relationships[0],  # First target only
                                    field_api=field_api,
                                    relationship_api=field.get("relationshipName"),
                                    cardinality="many_to_one",
                                    schema_hash=schema_hash,
                                )
                            )

                await self.repo.touch_heartbeat(org_id, token=token)
                await self.repo.replace_objects(org_id, schema_hash, objects_payload)
                await self.repo.replace_fields(org_id, schema_hash, fields_payload)
                await self.repo.replace_relationships(
                    org_id, schema_hash, relationships_payload
                )
                await self.repo.replace_terms(org_id, schema_hash, terms_payload)
                fls_hash = self._resolve_fls_hash(gateway)
                await self.repo.mark_quick_indexed(
                    org_id, schema_hash=schema_hash, fls_hash=fls_hash
                )
                await self.db.commit()
                emit(
                    "index.quick.ms",
                    org_id=org_id,
                    ms=int((time.perf_counter() - start) * 1000),
                    objects=len(objects_payload),
                    fields=len(fields_payload),
                )
                log.info(
                    "Quick schema index built for org %s: %s objects, %s fields",
                    org_id,
                    len(objects_payload),
                    len(fields_payload),
                )
                return schema_hash
        except RuntimeError:
            log.info(
                "Schema index quick build already running for %s, reusing existing hash",
                org_id,
            )
            refreshed = await self.repo.get_state(org_id)
            return refreshed.schema_hash if refreshed and refreshed.schema_hash else ""

    async def ensure_full_index(
        self,
        *,
        org_id: str,
        gateway: SalesforceGateway,
    ) -> Dict[str, int]:
        """Run full index including report usage, sampling, and score updates."""
        schema_hash = await self.ensure_quick_index(org_id=org_id, gateway=gateway)
        full_start = time.perf_counter()

        try:
            async with self._acquire_org_lock(org_id, stage="full") as token:
                checkpoint = await self.repo.get_checkpoint(org_id)
                start_offset = int(checkpoint.get("report_offset", 0) or 0)
                await self.repo.touch_heartbeat(org_id, token=token)

                usage_counts, cooccur, last_idx = await self._collect_report_usage(
                    org_id=org_id,
                    gateway=gateway,
                    start=start_offset,
                )
                await self.repo.save_report_usage(org_id, usage_counts, cooccur)
                await self.repo.set_checkpoint(org_id, {"report_offset": last_idx})
                await self.repo.touch_heartbeat(org_id, token=token)

                profiles = await self._profile_fields(
                    org_id=org_id,
                    gateway=gateway,
                    objects=("Opportunity", "Account", "Lead", "Contact"),
                )
                await self.repo.replace_field_profiles(org_id, profiles)
                await self.repo.touch_heartbeat(org_id, token=token)

                app_usage = await self.repo.app_usage(org_id)
                scores = await self._compute_field_scores(
                    org_id, usage_counts, app_usage, profiles, cooccur
                )
                await self.repo.update_term_weights(org_id, scores)

                await self.repo.mark_full_indexed(org_id, schema_hash)
                await self.db.commit()
                emit(
                    "index.full.ms",
                    org_id=org_id,
                    ms=int((time.perf_counter() - full_start) * 1000),
                    usage_fields=len(usage_counts),
                    profiled_objects=len(profiles),
                )
                return {
                    "schema_hash": schema_hash,
                    "usage_fields": len(usage_counts),
                    "profiled_objects": len(profiles),
                }
        except RuntimeError:
            log.info(
                "Schema index full build already running for %s; skipping duplicate call",
                org_id,
            )
            existing = await self.repo.get_state(org_id)
            return {
                "schema_hash": existing.schema_hash if existing else schema_hash,
                "usage_fields": 0,
                "profiled_objects": 0,
            }

    async def _collect_report_usage(
        self,
        *,
        org_id: str,
        gateway: SalesforceGateway,
        start: int = 0,
    ) -> Tuple[Counter, Dict[str, Counter], int]:
        limit = max(1, self.settings.AI_SCHEMA_INDEX_REPORTS_MAX)
        timer = time.perf_counter()
        query = (
            "SELECT Id, Name, DeveloperName, LastRunDate, LastReferencedDate "
            "FROM Report ORDER BY LastReferencedDate DESC NULLS LAST LIMIT {limit}{offset}"
        ).format(
            limit=limit,
            offset=f" OFFSET {start}" if start else "",
        )

        usage: Counter = Counter()
        cooccur: Dict[str, Counter] = defaultdict(Counter)

        try:
            result = await gateway.soql(query)
        except Exception as exc:  # noqa: BLE001
            log.warning("Failed to fetch report usage for %s: %s", org_id, exc)
            emit(
                "index.reports.ms",
                org_id=org_id,
                ms=int((time.perf_counter() - timer) * 1000),
                reports=0,
                error="report_fetch_failed",
            )
            return usage, cooccur, start

        for record in result.get("records", []) or []:
            report_id = record.get("Id")
            if not report_id:
                continue
            try:
                report_payload = await gateway.export_report(report_id)
            except Exception as exc:  # noqa: BLE001
                log.debug("Skipping report %s due to error: %s", report_id, exc)
                continue
            metadata = report_payload.get("reportMetadata", {}) or {}
            extended_info = report_payload.get("reportExtendedMetadata", {}) or {}
            column_info = {}
            for source in (
                extended_info.get("detailColumnInfo", {}),
                metadata.get("columnInfo", {}),
            ):
                if isinstance(source, dict):
                    column_info.update(source)

            def resolve_field(ref: str) -> Optional[str]:
                info = column_info.get(ref) or {}
                fq_name = info.get("fullyQualifiedName") or info.get("name")
                if fq_name and "." in fq_name:
                    return fq_name
                label = info.get("label")
                if label and "." in label:
                    return label
                return self._normalize_report_field(ref)

            columns = metadata.get("detailColumns", []) or []
            filters = [
                filt.get("column")
                for filt in metadata.get("reportFilters", []) or []
                if isinstance(filt, dict)
            ]
            touched_keys = [col for col in columns if col] + [
                col for col in filters if col
            ]
            resolved = {resolve_field(ref) for ref in touched_keys if ref}
            unique_fields = {ref for ref in resolved if ref}
            for field_ref in unique_fields:
                usage[field_ref] += 1
            for field_a in unique_fields:
                for field_b in unique_fields:
                    if field_a == field_b:
                        continue
                    cooccur[field_a][field_b] += 1  # Fixed: was outside loop
        emit(
            "index.reports.ms",
            org_id=org_id,
            ms=int((time.perf_counter() - timer) * 1000),
            reports=sum(usage.values()),
            offset=start,
        )
        last_index = start + len(result.get("records", []) or [])
        return usage, cooccur, last_index

    def _normalize_report_field(self, field_ref: str) -> str:
        """Normalize field reference from report metadata.

        Ensures consistent lowercase casing for custom field suffixes (__c, __r)
        to allow proper matching with schema fields.
        """
        if "." not in field_ref and "__" not in field_ref:
            return field_ref
        # Normalize custom suffixes to lowercase for consistent matching
        # Previous code was a no-op: .replace("__c", "__c") does nothing
        normalized = field_ref
        normalized = re.sub(r"__C\b", "__c", normalized)  # Uppercase __C to __c
        normalized = re.sub(r"__R\b", "__r", normalized)  # Uppercase __R to __r
        return normalized

    async def _profile_fields(
        self,
        *,
        org_id: str,
        gateway: SalesforceGateway,
        objects: Iterable[str],
    ) -> Dict[str, Dict[str, Dict[str, float]]]:
        profiles: Dict[str, Dict[str, Dict[str, float]]] = {}
        sample_limit = max(1, self.settings.AI_SCHEMA_INDEX_SAMPLE_LIMIT)
        for obj in objects:
            fields = await self.repo.queryable_fields(org_id, obj, limit=120)
            if not fields:
                continue
            select_clause = ", ".join(fields)
            soql = f"SELECT {select_clause} FROM {obj} LIMIT {sample_limit}"
            try:
                result = await gateway.soql(soql)
            except Exception as exc:  # noqa: BLE001
                log.debug("Sampling %s failed: %s", obj, exc)
                continue
            records = result.get("records", []) or []
            profiles[obj] = self._profile_records(fields, records)
        return profiles

    def _profile_records(
        self, fields: List[str], records: List[Dict[str, object]]
    ) -> Dict[str, Dict[str, float]]:
        field_stats: Dict[str, Dict[str, float]] = {}
        if not records:
            return field_stats
        total = len(records)
        for field_name in fields:
            normalized_values: List[str] = []
            non_null = 0
            for record in records:
                value = record.get(field_name)
                if value not in (None, ""):
                    non_null += 1
                    try:
                        normalized_values.append(
                            json.dumps(value, sort_keys=True, default=str)
                        )
                    except Exception:  # noqa: BLE001
                        normalized_values.append(str(value))
            distinct = len(set(normalized_values))
            top_values = []
            if normalized_values:
                counts = Counter(normalized_values)
                top_values = [val for val, _ in counts.most_common(5)]
            field_stats[field_name] = {
                "not_null_rate": non_null / total if total else 0.0,
                "distinct_ratio": distinct / total if total else 0.0,
                "top_values": top_values,
                "sample_values": [
                    normalized_values[idx]
                    for idx in range(min(5, len(normalized_values)))
                ]
                if normalized_values
                else [],
            }
        return field_stats

    async def _compute_field_scores(
        self,
        org_id: str,
        report_usage: Counter,
        app_usage: Dict[str, float],
        profiles: Dict[str, Dict[str, Dict[str, float]]],
        cooccur: Dict[str, Counter],
    ) -> Dict[str, float]:
        scores: Dict[str, float] = {}

        def norm(value: float) -> float:
            if value <= 0:
                return 0.0
            return min(1.0, value)

        field_refs = await self.repo.all_field_refs(org_id)
        for object_api, field_api, is_custom in field_refs:
            ref = f"{object_api}.{field_api}"
            usage_score = norm(
                report_usage.get(ref, 0)
                / max(1, self.settings.AI_SCHEMA_INDEX_REPORTS_MAX)
            )
            app_score = norm(app_usage.get(ref, 0) / 50.0)
            profile = profiles.get(object_api, {}).get(field_api, {})
            fill_score = norm(profile.get("not_null_rate", 0.0))
            standard_boost = 1.0 if not is_custom else 0.0
            cooccur_sum = float(sum(cooccur.get(ref, {}).values()))
            cooccur_score = norm(cooccur_sum / 25.0)
            score = (
                0.45 * usage_score
                + 0.25 * app_score
                + 0.15 * fill_score
                + 0.10 * standard_boost
                + 0.05 * cooccur_score
            )
            scores[ref] = round(score, 4)
        return scores
