from __future__ import annotations

import json
import logging
import os
import time
from datetime import datetime, timedelta, timezone
from collections import Counter
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple

from sqlalchemy import delete, func, select, or_, update, insert, text
from sqlalchemy.exc import DataError, ProgrammingError, IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession

from fmatch.saas.model_defs.salesforce_schema_index import (
    SalesforceSchemaField,
    SalesforceSchemaFieldAppUsage,
    SalesforceSchemaFieldProfile,
    SalesforceSchemaFieldUsage,
    SalesforceSchemaIndexState,
    SalesforceSchemaObject,
    SalesforceSchemaRelationship,
    SalesforceSchemaTerm,
)

logger = logging.getLogger(__name__)
EMBEDDING_RECHECK_SECONDS = int(os.getenv("SF_EMBEDDING_RECHECK_S", "300"))


class SalesforceSchemaIndexRepo:
    """Persistence helpers for the Salesforce schema index."""

    def __init__(self, db: AsyncSession) -> None:
        self.db = db
        self._embedding_column_available: Optional[bool] = None
        self._embedding_column_retry_at: Optional[float] = None
        self._embedding_recheck_seconds = EMBEDDING_RECHECK_SECONDS

    @staticmethod
    def _utcnow_naive() -> datetime:
        """Return a UTC datetime without tzinfo to match TIMESTAMP WITHOUT TIME ZONE columns."""
        return datetime.now(timezone.utc).replace(tzinfo=None)

    def _embedding_available(self) -> bool:
        if self._embedding_column_available is False:
            if (
                self._embedding_column_retry_at is not None
                and time.monotonic() >= self._embedding_column_retry_at
            ):
                self._embedding_column_available = None
                self._embedding_column_retry_at = None
                logger.info("Rechecking embedding column after cooldown")
            else:
                return False
        return True

    @staticmethod
    def _as_naive(dt: Optional[datetime]) -> Optional[datetime]:
        if dt is None:
            return None
        if dt.tzinfo:
            return dt.replace(tzinfo=None)
        return dt

    # ------------------------------------------------------------------ #
    # Index state helpers
    # ------------------------------------------------------------------ #
    async def get_state(self, org_id: str) -> Optional[SalesforceSchemaIndexState]:
        return await self.db.get(SalesforceSchemaIndexState, org_id)

    async def get_status_row(self, org_id: str) -> Optional[SalesforceSchemaIndexState]:
        return await self.get_state(org_id)

    async def mark_in_progress(
        self, org_id: str, *, stage: str, token: str
    ) -> SalesforceSchemaIndexState:
        now = self._utcnow_naive()
        stale_cutoff = now - timedelta(minutes=5)

        update_stmt = (
            update(SalesforceSchemaIndexState)
            .where(SalesforceSchemaIndexState.org_id == org_id)
            .where(
                or_(
                    SalesforceSchemaIndexState.in_progress == False,
                    SalesforceSchemaIndexState.heartbeat_at.is_(None),
                    SalesforceSchemaIndexState.heartbeat_at < stale_cutoff,
                )
            )
            .values(
                status=f"{stage}_running",
                in_progress=True,
                stage=stage,
                runner_token=token,
                heartbeat_at=now,
                updated_at=now,
            )
        )
        update_result = await self.db.execute(update_stmt)
        if (update_result.rowcount or 0) == 1:
            row = await self.get_state(org_id)
            if row:
                return row

        row = SalesforceSchemaIndexState(
            org_id=org_id,
            schema_hash=None,
            status=f"{stage}_running",
            in_progress=True,
            stage=stage,
            runner_token=token,
            heartbeat_at=now,
            checkpoint_json="{}",
            updated_at=now,
        )
        self.db.add(row)
        try:
            await self.db.flush()
            return row
        except IntegrityError as exc:
            await self.db.rollback()
            raise RuntimeError("Index already running") from exc

    async def clear_in_progress(self, org_id: str, *, token: Optional[str]) -> None:
        row = await self.get_state(org_id)
        if not row:
            return
        if token and row.runner_token and row.runner_token != token:
            return
        row.in_progress = False
        row.stage = None
        row.runner_token = None
        row.heartbeat_at = None
        if row.status.endswith("_running"):
            row.status = "idle"
        await self.db.flush()

    async def touch_heartbeat(self, org_id: str, *, token: Optional[str]) -> None:
        # Token is required to prevent heartbeat hijacking
        if not token:
            raise ValueError("Token required for heartbeat - prevents unauthorized lock takeover")
        stmt = (
            update(SalesforceSchemaIndexState)
            .where(
                SalesforceSchemaIndexState.org_id == org_id,
                SalesforceSchemaIndexState.runner_token == token,
            )
            .values(heartbeat_at=self._utcnow_naive())
        )
        await self.db.execute(stmt)

    async def get_checkpoint(self, org_id: str) -> Dict[str, Any]:
        row = await self.get_state(org_id)
        if not row or not row.checkpoint_json:
            return {}
        try:
            return json.loads(row.checkpoint_json)
        except Exception:  # noqa: BLE001
            return {}

    async def set_checkpoint(
        self, org_id: str, checkpoint: Optional[Dict[str, Any]]
    ) -> None:
        row = await self.get_state(org_id)
        payload = json.dumps(checkpoint or {}, ensure_ascii=False)
        if not row:
            row = SalesforceSchemaIndexState(
                org_id=org_id,
                checkpoint_json=payload,
            )
            self.db.add(row)
        else:
            row.checkpoint_json = payload
        await self.db.flush()

    async def mark_quick_indexed(
        self,
        org_id: str,
        *,
        schema_hash: Optional[str],
        fls_hash: Optional[str] = None,
    ) -> None:
        now = self._utcnow_naive()
        row = await self.get_state(org_id)
        if not row:
            row = SalesforceSchemaIndexState(
                org_id=org_id,
                schema_hash=schema_hash,
                quick_indexed_at=now,
                status="ready",
                fls_hash=fls_hash,
            )
            self.db.add(row)
        else:
            row.schema_hash = schema_hash
            row.quick_indexed_at = now
            row.status = "ready"
            if fls_hash:
                row.fls_hash = fls_hash
        row.in_progress = False
        row.stage = None
        row.runner_token = None
        row.heartbeat_at = None
        await self.db.flush()

    async def upsert_state(
        self,
        org_id: str,
        *,
        schema_hash: Optional[str],
        status: str,
        quick_indexed_at: Optional[datetime] = None,
        full_indexed_at: Optional[datetime] = None,
        error_json: Optional[str] = None,
    ) -> SalesforceSchemaIndexState:
        quick_indexed_at = self._as_naive(quick_indexed_at)
        full_indexed_at = self._as_naive(full_indexed_at)
        row = await self.db.get(SalesforceSchemaIndexState, org_id)
        if not row:
            row = SalesforceSchemaIndexState(
                org_id=org_id,
                schema_hash=schema_hash,
                status=status,
                quick_indexed_at=quick_indexed_at,
                full_indexed_at=full_indexed_at,
                error_json=error_json,
            )
            self.db.add(row)
        else:
            row.schema_hash = schema_hash
            row.status = status
            if quick_indexed_at is not None:
                row.quick_indexed_at = quick_indexed_at
            if full_indexed_at is not None:
                row.full_indexed_at = full_indexed_at
            row.error_json = error_json
            row.updated_at = self._utcnow_naive()
            if not row.checkpoint_json:
                row.checkpoint_json = "{}"
        await self.db.flush()
        return row

    # ------------------------------------------------------------------ #
    # Bulk persistence helpers
    # ------------------------------------------------------------------ #
    async def replace_objects(
        self, org_id: str, schema_hash: str, rows: Iterable[SalesforceSchemaObject]
    ) -> None:
        rows = list(rows)
        if not self._embedding_available():
            await self._replace_objects_without_embedding(org_id, rows)
            return
        try:
            async with self.db.begin_nested():
                await self.db.execute(
                    delete(SalesforceSchemaObject).where(
                        SalesforceSchemaObject.org_id == org_id
                    )
                )
                self.db.add_all(rows)
                await self.db.flush()
        except ProgrammingError as exc:
            # Savepoint auto-rollback handled by begin_nested() context manager
            if not self._handle_missing_embedding(exc, stage="objects"):
                raise
            await self._replace_objects_without_embedding(org_id, rows)
        except DataError as exc:
            # Savepoint auto-rollback handled by begin_nested() context manager
            logger.error("Data error in replace_objects for %s: %s", org_id, exc)
            raise

    async def replace_fields(
        self,
        org_id: str,
        schema_hash: str,
        rows: Iterable[SalesforceSchemaField],
    ) -> None:
        rows = list(rows)
        if not self._embedding_available():
            await self._replace_fields_without_embedding(org_id, rows)
            return
        try:
            async with self.db.begin_nested():
                await self.db.execute(
                    delete(SalesforceSchemaField).where(
                        SalesforceSchemaField.org_id == org_id
                    )
                )
                self.db.add_all(rows)
                await self.db.flush()
        except ProgrammingError as exc:
            # Savepoint auto-rollback handled by begin_nested() context manager
            if not self._handle_missing_embedding(exc, stage="fields"):
                raise
            await self._replace_fields_without_embedding(org_id, rows)
        except DataError as exc:
            # Savepoint auto-rollback handled by begin_nested() context manager
            logger.error("Data error in replace_fields for %s: %s", org_id, exc)
            raise

    async def replace_relationships(
        self,
        org_id: str,
        schema_hash: str,
        rows: Iterable[SalesforceSchemaRelationship],
    ) -> None:
        try:
            async with self.db.begin_nested():
                await self.db.execute(
                    delete(SalesforceSchemaRelationship).where(
                        SalesforceSchemaRelationship.org_id == org_id
                    )
                )
                self.db.add_all(list(rows))
                await self.db.flush()
        except DataError as exc:
            # Savepoint auto-rollback handled by begin_nested() context manager
            logger.error("Data error in replace_relationships for %s: %s", org_id, exc)
            raise

    async def replace_terms(
        self,
        org_id: str,
        schema_hash: str,
        rows: Iterable[SalesforceSchemaTerm],
    ) -> None:
        try:
            async with self.db.begin_nested():
                await self.db.execute(
                    delete(SalesforceSchemaTerm).where(SalesforceSchemaTerm.org_id == org_id)
                )
                self.db.add_all(list(rows))
                await self.db.flush()
        except DataError as exc:
            # Savepoint auto-rollback handled by begin_nested() context manager
            logger.error("Data error in replace_terms for %s: %s", org_id, exc)
            raise

    # Usage / profile upserters
    async def upsert_field_usage(
        self,
        *,
        org_id: str,
        field_ref: str,
        report_count: int,
        last_run_at: Optional[datetime],
        last_referenced_at: Optional[datetime],
        schema_hash: Optional[str],
    ) -> None:
        stmt = select(SalesforceSchemaFieldUsage).where(
            SalesforceSchemaFieldUsage.org_id == org_id,
            SalesforceSchemaFieldUsage.field_ref == field_ref,
        )
        existing = (await self.db.execute(stmt)).scalar_one_or_none()
        if existing:
            existing.report_count = report_count
            existing.last_run_at = last_run_at
            existing.last_referenced_at = last_referenced_at
            existing.schema_hash = schema_hash
            existing.indexed_at = self._utcnow_naive()
        else:
            row = SalesforceSchemaFieldUsage(
                org_id=org_id,
                field_ref=field_ref,
                report_count=report_count,
                last_run_at=last_run_at,
                last_referenced_at=last_referenced_at,
                schema_hash=schema_hash,
            )
            self.db.add(row)
        await self.db.flush()

    async def upsert_field_profile(
        self,
        *,
        org_id: str,
        field_ref: str,
        not_null_rate: Optional[float],
        distinct_ratio: Optional[float],
        top_values: Optional[Sequence[str]] = None,
        sample_values: Optional[Sequence[str]] = None,
    ) -> None:
        stmt = select(SalesforceSchemaFieldProfile).where(
            SalesforceSchemaFieldProfile.org_id == org_id,
            SalesforceSchemaFieldProfile.field_ref == field_ref,
        )
        existing = (await self.db.execute(stmt)).scalar_one_or_none()
        payload = json.dumps(list(top_values or []), ensure_ascii=False)
        sample = json.dumps(list(sample_values or []), ensure_ascii=False)
        if existing:
            existing.not_null_rate = not_null_rate
            existing.distinct_ratio = distinct_ratio
            existing.top_values_json = payload
            existing.sample_values_json = sample
            existing.updated_at = self._utcnow_naive()
        else:
            row = SalesforceSchemaFieldProfile(
                org_id=org_id,
                field_ref=field_ref,
                not_null_rate=not_null_rate,
                distinct_ratio=distinct_ratio,
                top_values_json=payload,
                sample_values_json=sample,
            )
            self.db.add(row)
        await self.db.flush()

    async def upsert_app_usage(
        self,
        *,
        org_id: str,
        field_ref: str,
        query_delta: int,
        preview_delta: int,
        cooccur: Optional[dict] = None,
    ) -> None:
        stmt = select(SalesforceSchemaFieldAppUsage).where(
            SalesforceSchemaFieldAppUsage.org_id == org_id,
            SalesforceSchemaFieldAppUsage.field_ref == field_ref,
        )
        existing = (await self.db.execute(stmt)).scalar_one_or_none()
        if existing:
            existing.query_count_30d = max(0, existing.query_count_30d + query_delta)
            existing.preview_count_30d = max(
                0, existing.preview_count_30d + preview_delta
            )
            if cooccur:
                current = json.loads(existing.cooccur_json or "{}")
                for k, v in cooccur.items():
                    current[k] = current.get(k, 0) + v
                existing.cooccur_json = json.dumps(current, ensure_ascii=False)
            existing.updated_at = self._utcnow_naive()
        else:
            row = SalesforceSchemaFieldAppUsage(
                org_id=org_id,
                field_ref=field_ref,
                query_count_30d=max(0, query_delta),
                preview_count_30d=max(0, preview_delta),
                cooccur_json=json.dumps(cooccur or {}, ensure_ascii=False),
            )
            self.db.add(row)
        await self.db.flush()

    # ------------------------------------------------------------------ #
    # Query helpers for discover/planner
    # ------------------------------------------------------------------ #
    async def list_objects(
        self, org_id: str, limit: int = 50
    ) -> List[SalesforceSchemaObject]:
        stmt = (
            select(SalesforceSchemaObject)
            .where(SalesforceSchemaObject.org_id == org_id)
            .order_by(
                SalesforceSchemaObject.is_custom.asc(),
                SalesforceSchemaObject.label.asc(),
            )
            .limit(limit)
        )
        return list((await self.db.execute(stmt)).scalars().all())

    async def list_fields(
        self, org_id: str, object_api: str
    ) -> List[SalesforceSchemaField]:
        stmt = (
            select(SalesforceSchemaField)
            .where(
                SalesforceSchemaField.org_id == org_id,
                SalesforceSchemaField.object_api == object_api,
            )
            .order_by(SalesforceSchemaField.field_api.asc())
        )
        return list((await self.db.execute(stmt)).scalars().all())

    async def top_terms(
        self, org_id: str, term: str, limit: int = 10
    ) -> List[Tuple[str, float]]:
        stmt = (
            select(SalesforceSchemaTerm.field_ref, SalesforceSchemaTerm.weight)
            .where(
                SalesforceSchemaTerm.org_id == org_id,
                SalesforceSchemaTerm.term == term.lower(),
            )
            .order_by(SalesforceSchemaTerm.weight.desc())
            .limit(limit)
        )
        return list((await self.db.execute(stmt)).all())

    async def replace_field_usage(self, org_id: str, usage: Counter) -> None:
        try:
            async with self.db.begin_nested():
                await self.db.execute(
                    delete(SalesforceSchemaFieldUsage).where(
                        SalesforceSchemaFieldUsage.org_id == org_id
                    )
                )
                now = self._utcnow_naive()
                rows = [
                    SalesforceSchemaFieldUsage(
                        org_id=org_id,
                        field_ref=field_ref,
                        report_count=int(count),
                        indexed_at=now,
                    )
                    for field_ref, count in usage.items()
                ]
                if rows:
                    self.db.add_all(rows)
                await self.db.flush()
        except DataError as exc:
            # Savepoint auto-rollback handled by begin_nested() context manager
            logger.error("Data error in replace_field_usage for %s: %s", org_id, exc)
            raise

    async def save_report_usage(
        self, org_id: str, usage: Counter, cooccur: Dict[str, Counter]
    ) -> None:
        await self.replace_field_usage(org_id, usage)
        # co-occurrence matrix is kept in-memory for scoring; persist in future if required

    async def replace_field_profiles(
        self, org_id: str, profiles: Dict[str, Dict[str, Dict[str, float]]]
    ) -> None:
        now = self._utcnow_naive()
        rows = []
        for obj, field_map in profiles.items():
            for field_api, stats in field_map.items():
                field_ref = f"{obj}.{field_api}"
                rows.append(
                    SalesforceSchemaFieldProfile(
                        org_id=org_id,
                        field_ref=field_ref,
                        not_null_rate=stats.get("not_null_rate"),
                        distinct_ratio=stats.get("distinct_ratio"),
                        top_values_json=json.dumps(
                            stats.get("top_values", []) or [], ensure_ascii=False
                        ),
                        sample_values_json=json.dumps(
                            stats.get("sample_values", []) or [], ensure_ascii=False
                        ),
                        updated_at=now,
                    )
                )
        try:
            async with self.db.begin_nested():
                await self.db.execute(
                    delete(SalesforceSchemaFieldProfile).where(
                        SalesforceSchemaFieldProfile.org_id == org_id
                    )
                )
                if rows:
                    self.db.add_all(rows)
                await self.db.flush()
        except DataError as exc:
            # Savepoint auto-rollback handled by begin_nested() context manager
            logger.error("Data error in replace_field_profiles for %s: %s", org_id, exc)
            raise

    async def update_term_weights(self, org_id: str, scores: Dict[str, float]) -> None:
        if not scores:
            return
        stmt = select(SalesforceSchemaTerm).where(SalesforceSchemaTerm.org_id == org_id)
        rows = (await self.db.execute(stmt)).scalars().all()
        now = self._utcnow_naive()
        for row in rows:
            score = scores.get(row.field_ref)
            if score is None:
                continue
            row.weight = float(max(score, row.weight or 0.0))
            row.indexed_at = now
        await self.db.flush()

    async def app_usage(self, org_id: str) -> Dict[str, float]:
        stmt = select(
            SalesforceSchemaFieldAppUsage.field_ref,
            SalesforceSchemaFieldAppUsage.query_count_30d,
            SalesforceSchemaFieldAppUsage.preview_count_30d,
        ).where(SalesforceSchemaFieldAppUsage.org_id == org_id)
        return {
            field_ref: float((query or 0) + (preview or 0))
            for field_ref, query, preview in await self.db.execute(stmt)
        }

    async def cooccur_sum(self, org_id: str, field_ref: str) -> float:
        stmt = select(SalesforceSchemaFieldAppUsage.cooccur_json).where(
            SalesforceSchemaFieldAppUsage.org_id == org_id,
            SalesforceSchemaFieldAppUsage.field_ref == field_ref,
        )
        raw = (await self.db.execute(stmt)).scalar_one_or_none()
        if not raw:
            return 0.0
        try:
            data = json.loads(raw)
        except Exception:
            return 0.0
        if isinstance(data, dict):
            return float(sum(data.values()))
        return 0.0

    async def all_field_refs(self, org_id: str) -> List[Tuple[str, str, bool]]:
        stmt = select(
            SalesforceSchemaField.object_api,
            SalesforceSchemaField.field_api,
            SalesforceSchemaField.is_custom,
        ).where(SalesforceSchemaField.org_id == org_id)
        return [(row[0], row[1], bool(row[2])) for row in await self.db.execute(stmt)]

    async def queryable_fields(
        self, org_id: str, object_api: str, limit: int = 120
    ) -> List[str]:
        stmt = (
            select(SalesforceSchemaField.field_api)
            .where(
                SalesforceSchemaField.org_id == org_id,
                SalesforceSchemaField.object_api == object_api,
                SalesforceSchemaField.queryable.is_(True),
            )
            .limit(limit)
        )
        return [row[0] for row in await self.db.execute(stmt)]

    # ------------------------------------------------------------------ #
    # Semantic search helpers (PGVector)
    # ------------------------------------------------------------------ #
    async def find_similar_objects(
        self, org_id: str, vector: List[float], limit: int = 5
    ) -> List[SalesforceSchemaObject]:
        if not vector:
            return []
        if not self._embedding_available():
            return []
        # Use raw SQL since embedding column is not in ORM model
        vector_str = "[" + ",".join(str(v) for v in vector) + "]"
        sql = text("""
            SELECT org_id, object_api FROM sf_schema_objects
            WHERE org_id = :org_id AND embedding IS NOT NULL
            ORDER BY embedding <=> CAST(:vector AS vector)
            LIMIT :limit
        """)
        try:
            result = await self.db.execute(
                sql, {"org_id": org_id, "vector": vector_str, "limit": limit}
            )
            object_apis = [row[1] for row in result.fetchall()]
            if not object_apis:
                return []
            # Fetch full objects by their APIs
            stmt = select(SalesforceSchemaObject).where(
                SalesforceSchemaObject.org_id == org_id,
                SalesforceSchemaObject.object_api.in_(object_apis),
            )
            return list((await self.db.execute(stmt)).scalars().all())
        except ProgrammingError as exc:
            if self._handle_missing_embedding(exc, stage="objects"):
                return []
            raise

    async def find_similar_fields(
        self,
        org_id: str,
        vector: List[float],
        limit: int = 20,
        object_filter: Optional[List[str]] = None,
    ) -> List[SalesforceSchemaField]:
        if not vector:
            return []
        if not self._embedding_available():
            return []
        # Use raw SQL since embedding column is not in ORM model
        vector_str = "[" + ",".join(str(v) for v in vector) + "]"
        if object_filter:
            sql = text("""
                SELECT org_id, object_api, field_api FROM sf_schema_fields
                WHERE org_id = :org_id AND embedding IS NOT NULL
                  AND object_api = ANY(:object_filter)
                ORDER BY embedding <=> CAST(:vector AS vector)
                LIMIT :limit
            """)
            params = {
                "org_id": org_id,
                "vector": vector_str,
                "limit": limit,
                "object_filter": object_filter,
            }
        else:
            sql = text("""
                SELECT org_id, object_api, field_api FROM sf_schema_fields
                WHERE org_id = :org_id AND embedding IS NOT NULL
                ORDER BY embedding <=> CAST(:vector AS vector)
                LIMIT :limit
            """)
            params = {"org_id": org_id, "vector": vector_str, "limit": limit}
        try:
            result = await self.db.execute(sql, params)
            rows = result.fetchall()
            if not rows:
                return []
            # Fetch full fields by their keys
            conditions = [
                (SalesforceSchemaField.org_id == org_id)
                & (SalesforceSchemaField.object_api == r[1])
                & (SalesforceSchemaField.field_api == r[2])
                for r in rows
            ]
            stmt = select(SalesforceSchemaField).where(or_(*conditions))
            return list((await self.db.execute(stmt)).scalars().all())
        except ProgrammingError as exc:
            if self._handle_missing_embedding(exc, stage="fields"):
                return []
            raise

    def _handle_missing_embedding(self, exc: Exception, *, stage: str) -> bool:
        message = str(exc).lower()
        if "embedding" in message and "does not exist" in message:
            if self._embedding_column_available is not False:
                self._embedding_column_available = False
                if self._embedding_recheck_seconds > 0:
                    self._embedding_column_retry_at = (
                        time.monotonic() + self._embedding_recheck_seconds
                    )
                logger.warning(
                    "Schema index embeddings disabled (%s): missing embedding column. "
                    "Run migrations to add pgvector columns.",
                    stage,
                )
            return True
        return False

    async def _replace_objects_without_embedding(
        self, org_id: str, rows: List[SalesforceSchemaObject]
    ) -> None:
        async with self.db.begin_nested():
            await self.db.execute(
                delete(SalesforceSchemaObject).where(
                    SalesforceSchemaObject.org_id == org_id
                )
            )
            # embedding column exists in DB but not in ORM model
            payload = [row.model_dump() for row in rows]
            if payload:
                await self.db.execute(
                    insert(SalesforceSchemaObject.__table__).values(payload)
                )
            await self.db.flush()

    async def _replace_fields_without_embedding(
        self, org_id: str, rows: List[SalesforceSchemaField]
    ) -> None:
        async with self.db.begin_nested():
            await self.db.execute(
                delete(SalesforceSchemaField).where(
                    SalesforceSchemaField.org_id == org_id
                )
            )
            # embedding column exists in DB but not in ORM model
            payload = [row.model_dump() for row in rows]
            if payload:
                await self.db.execute(
                    insert(SalesforceSchemaField.__table__).values(payload)
                )
            await self.db.flush()

    async def rank_for_terms(
        self,
        *,
        org_id: str,
        terms: Sequence[str],
        limit: int = 20,
    ) -> List[Tuple[str, float, str, str, str, int, float]]:
        normalized_terms = [term.lower() for term in terms if term]
        seen: Dict[str, float] = {}
        results: List[Tuple[str, float, str, str, str, int, float]] = []

        for term in normalized_terms or ["amount"]:
            term_rows = (
                await self.db.execute(
                    select(
                        SalesforceSchemaTerm.field_ref,
                        SalesforceSchemaTerm.weight,
                    )
                    .where(
                        SalesforceSchemaTerm.org_id == org_id,
                        SalesforceSchemaTerm.term == term,
                    )
                    .order_by(SalesforceSchemaTerm.weight.desc())
                    .limit(limit)
                )
            ).all()
            if not term_rows:
                continue

            field_refs = [row[0] for row in term_rows if row[0]]
            pairs: List[Tuple[str, str]] = []
            for ref in field_refs:
                if "." in ref:
                    obj, fld = ref.split(".", 1)
                    pairs.append((obj, fld))

            field_meta: Dict[Tuple[str, str], Dict[str, Any]] = {}
            if pairs:
                conditions = [
                    (SalesforceSchemaField.object_api == obj)
                    & (SalesforceSchemaField.field_api == fld)
                    for obj, fld in pairs
                ]
                where_clause = (
                    or_(*conditions) if len(conditions) > 1 else conditions[0]
                )
                field_rows = (
                    await self.db.execute(
                        select(
                            SalesforceSchemaField.object_api,
                            SalesforceSchemaField.field_api,
                            SalesforceSchemaField.label,
                        ).where(
                            SalesforceSchemaField.org_id == org_id,
                            where_clause,
                        )
                    )
                ).all()
                for obj, fld, label in field_rows:
                    field_meta[(obj, fld)] = {
                        "label": label,
                        "object_api": obj,
                        "field_api": fld,
                    }

            if field_refs:
                usage_rows = (
                    await self.db.execute(
                        select(
                            SalesforceSchemaFieldUsage.field_ref,
                            SalesforceSchemaFieldUsage.report_count,
                        ).where(
                            SalesforceSchemaFieldUsage.org_id == org_id,
                            SalesforceSchemaFieldUsage.field_ref.in_(field_refs),
                        )
                    )
                ).all()
                profile_rows = (
                    await self.db.execute(
                        select(
                            SalesforceSchemaFieldProfile.field_ref,
                            SalesforceSchemaFieldProfile.not_null_rate,
                        ).where(
                            SalesforceSchemaFieldProfile.org_id == org_id,
                            SalesforceSchemaFieldProfile.field_ref.in_(field_refs),
                        )
                    )
                ).all()
            else:
                usage_rows = []
                profile_rows = []
            usage_map = {ref: count for ref, count in usage_rows}
            profile_map = {ref: rate for ref, rate in profile_rows}

            for field_ref, weight in term_rows:
                if field_ref not in seen or weight > seen[field_ref]:
                    obj, fld = (
                        field_ref.split(".", 1) if "." in field_ref else ("", field_ref)
                    )
                    meta = field_meta.get(
                        (obj, fld), {"label": fld, "object_api": obj, "field_api": fld}
                    )
                    seen[field_ref] = weight
                    results.append(
                        (
                            field_ref,
                            float(weight),
                            meta.get("label") or meta.get("field_api"),
                            meta.get("object_api"),
                            meta.get("field_api"),
                            int(usage_map.get(field_ref, 0) or 0),
                            float(profile_map.get(field_ref, 0.0) or 0.0),
                        )
                    )
        # Deduplicate while preserving the highest weight
        unique = {}
        for item in results:
            if item[0] not in unique or item[1] > unique[item[0]][1]:
                unique[item[0]] = item
        ranked = sorted(unique.values(), key=lambda row: row[1], reverse=True)
        return ranked[:limit]

    async def mark_full_indexed(self, org_id: str, schema_hash: Optional[str]) -> None:
        await self.upsert_state(
            org_id,
            schema_hash=schema_hash,
            status="ready",
            full_indexed_at=self._utcnow_naive(),
        )
        await self.clear_in_progress(org_id, token=None)
        await self.set_checkpoint(org_id, {})

    async def status(self, org_id: str, *, ttl_hours: int = 24) -> dict:
        obj_count = (
            await self.db.execute(
                select(func.count())
                .select_from(SalesforceSchemaObject)
                .where(SalesforceSchemaObject.org_id == org_id)
            )
        ).scalar_one()
        field_count = (
            await self.db.execute(
                select(func.count())
                .select_from(SalesforceSchemaField)
                .where(SalesforceSchemaField.org_id == org_id)
            )
        ).scalar_one()
        state = await self.get_state(org_id)
        now = self._utcnow_naive()
        ttl = timedelta(hours=max(1, ttl_hours))
        full_dt = self._as_naive(state.full_indexed_at) if state else None
        quick_dt = self._as_naive(state.quick_indexed_at) if state else None
        if state and state.in_progress:
            index_status = "building"
        elif state and full_dt and (now - full_dt) <= ttl:
            index_status = "full"
        elif state and quick_dt and (now - quick_dt) <= ttl:
            index_status = "quick"
        else:
            index_status = "stale"
        source = "index" if index_status in {"full", "quick"} else "live"
        return {
            "org_id": org_id,
            "objects": obj_count,
            "fields": field_count,
            "schema_hash": state.schema_hash if state else None,
            "quick_indexed_at": state.quick_indexed_at.isoformat()
            if state and state.quick_indexed_at
            else None,
            "full_indexed_at": state.full_indexed_at.isoformat()
            if state and state.full_indexed_at
            else None,
            "status": state.status if state else "unknown",
            "in_progress": bool(state.in_progress) if state else False,
            "stage": state.stage if state else None,
            "index_status": index_status,
            "index_source": source,
            "fls_hash": state.fls_hash if state else None,
        }

    async def field_usage_score(self, org_id: str, field_ref: str) -> float:
        """Compute a basic usage score across report + app usage."""
        report_stmt = select(SalesforceSchemaFieldUsage.report_count).where(
            SalesforceSchemaFieldUsage.org_id == org_id,
            SalesforceSchemaFieldUsage.field_ref == field_ref,
        )
        report_count = (await self.db.execute(report_stmt)).scalar_one_or_none() or 0

        app_stmt = select(
            SalesforceSchemaFieldAppUsage.query_count_30d
            + SalesforceSchemaFieldAppUsage.preview_count_30d
        ).where(
            SalesforceSchemaFieldAppUsage.org_id == org_id,
            SalesforceSchemaFieldAppUsage.field_ref == field_ref,
        )
        app_count = (await self.db.execute(app_stmt)).scalar_one_or_none() or 0
        return float(report_count) * 0.7 + float(app_count) * 0.3

    async def field_profile(
        self, org_id: str, field_ref: str
    ) -> Optional[SalesforceSchemaFieldProfile]:
        stmt = select(SalesforceSchemaFieldProfile).where(
            SalesforceSchemaFieldProfile.org_id == org_id,
            SalesforceSchemaFieldProfile.field_ref == field_ref,
        )
        return (await self.db.execute(stmt)).scalar_one_or_none()

    async def search_fields(
        self, org_id: str, term: str, limit: int = 10
    ) -> List[SalesforceSchemaField]:
        """Search by term fallback (like contains)."""
        like_term = f"%{term.lower()}%"
        stmt = (
            select(SalesforceSchemaField)
            .where(
                SalesforceSchemaField.org_id == org_id,
                func.lower(SalesforceSchemaField.label).like(like_term)
                | func.lower(SalesforceSchemaField.field_api).like(like_term),
            )
            .limit(limit)
        )
        return list((await self.db.execute(stmt)).scalars().all())

    async def list_relationships(
        self, org_id: str, from_object: str
    ) -> List[SalesforceSchemaRelationship]:
        stmt = (
            select(SalesforceSchemaRelationship)
            .where(
                SalesforceSchemaRelationship.org_id == org_id,
                SalesforceSchemaRelationship.from_object == from_object,
            )
            .order_by(SalesforceSchemaRelationship.relationship_api.asc())
        )
        return list((await self.db.execute(stmt)).scalars().all())
