from __future__ import annotations

import asyncio
from collections import defaultdict
import hashlib
import json
import logging
import os
import re
import threading
import time
from copy import deepcopy
from dataclasses import dataclass, field
from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union


from sqlalchemy.ext.asyncio import AsyncSession
from pydantic import (
    BaseModel,
    ConfigDict,
    Field,
    ValidationError,
    field_validator,
    model_validator,
)

from fmatch.saas.settings import (
    AI_REPORTS_ENABLED,
    AI_REPORTS_ENTERPRISE_MODEL,
    AI_REPORTS_ALLOW_FIELDS,
    AI_REPORTS_ALLOW_OBJECTS,
    AI_REPORTS_DENY_OBJECTS,
    AI_REPORTS_MAX_FIELDS,
    AI_REPORTS_MAX_OBJECTS,
    AI_REPORTS_CRITIC_MODEL,
    AI_REPORTS_PLANNER_MODEL,
    AI_REPORTS_PREVIEW_LIMIT,
    AI_REPORTS_PRO_MODEL,
    AI_REPORTS_RESEARCHER_MODEL,
    AI_REPORTS_RESEARCH_CACHE_TTL,
    AI_REPORTS_RESEARCHER_MAX_OUTPUT_TOKENS,
    AI_REPORTS_RESEARCHER_TIMEOUT_S,
    AI_REPORTS_PLANNER_MAX_OUTPUT_TOKENS,
    AI_REPORTS_PLANNER_TIMEOUT_S,
    AI_REPORTS_ROW_CAPS,
    AI_REPORTS_SCHEMA_CACHE_TTL,
    AI_REPORTS_SKIP_CRITIC_SIMPLE,
    AI_REPORTS_TIMEOUT_S,
    AI_REPORTS_THINKING_LEVELS,
    AI_REPORTS_VALIDATOR_MODEL,
    AI_SCHEMA_INDEX_TTL_HOURS,
    VERTEX_LOCATION,
    VERTEX_PROJECT,
)

from fmatch.saas.db import AsyncSessionLocal
from fmatch.saas.repos.salesforce_schema_index_repo import SalesforceSchemaIndexRepo
from fmatch.saas.services.telemetry import emit

from .cache import HybridCache, make_cache_key
from fmatch.saas.providers.vertex_ai_provider import VertexAIProvider
from .salesforce_gateway import SalesforceGateway
from .salesforce_schema_index import SalesforceSchemaIndexer

try:  # pragma: no cover - optional dependency
    import vertexai
    from vertexai.language_models import TextEmbeddingModel
except Exception:  # noqa: BLE001 - optional dependency
    vertexai = None  # type: ignore
    TextEmbeddingModel = None  # type: ignore


log = logging.getLogger(__name__)


class PlannerUnavailable(RuntimeError):
    """Raised when no LLM planner is configured or reachable."""


class PlanValidationError(ValueError):
    """Raised when the AI generated plan fails schema validation."""


class SOQLRepairError(ValueError):
    """Raised when SOQL execution fails with a potentially repairable error."""

    def __init__(
        self,
        message: str,
        error_type: str,
        invalid_fields: Optional[List[str]] = None,
        original_error: Optional[str] = None,
    ):
        super().__init__(message)
        self.error_type = error_type
        self.invalid_fields = invalid_fields or []
        self.original_error = original_error


# Repair loop constants
MAX_REPAIR_ATTEMPTS = 2

# Error taxonomy for SOQL repair decisions
SOQL_ERROR_TAXONOMY = {
    # Fixable by removing/correcting fields
    "FIXABLE_FIELD": [
        "INVALID_FIELD",
        "No such column",
        "invalid field",
        "doesn't exist",
        "unknown field",
    ],
    # Fixable by correcting relationship names
    "FIXABLE_RELATIONSHIP": [
        "Didn't understand relationship",
        "invalid relationship",
        "Relationship '.+' does not exist",
        "child relationship",
    ],
    # Fixable by correcting syntax
    "FIXABLE_SYNTAX": [
        "MALFORMED_QUERY",
        "unexpected token",
        "syntax error",
        "invalid query",
        "missing operator",
    ],
    # Fixable by adjusting aggregation
    "FIXABLE_AGGREGATION": [
        "grouped",
        "aggregate",
        "must be grouped",
        "field must be grouped",
        "GROUP BY",
    ],
    # Not fixable - permission/limit/timeout issues
    "NOT_FIXABLE": [
        "INSUFFICIENT_ACCESS",
        "permission",
        "QUERY_TIMEOUT",
        "REQUEST_LIMIT_EXCEEDED",
        "API_LIMIT_EXCEEDED",
        "RATE_LIMIT",
        "quota",
    ],
}


def _repair_json(malformed: str) -> Optional[str]:
    """Attempt to repair common JSON issues from LLM output.

    Handles:
    - Trailing commas before ] or }
    - Unquoted keys
    - Single quotes instead of double quotes
    - Truncated JSON (attempts bracket balancing)
    - Control characters

    Returns:
        Repaired JSON string if successful, None if repair fails.
    """
    if not malformed or not malformed.strip():
        return None

    text = malformed.strip()

    # Remove control characters except newline/tab
    text = re.sub(r"[\x00-\x08\x0b\x0c\x0e-\x1f]", "", text)

    # Fix single quotes to double quotes (careful with apostrophes in text)
    # Only replace quotes that look like JSON string delimiters
    # Pattern: 'key': or : 'value' or ['item']
    text = re.sub(r"(?<=[\[{,:])\s*'([^']*)'", r' "\1"', text)
    text = re.sub(r"'([^']*)'(?=\s*[:,\]}])", r'"\1"', text)

    # Fix unquoted keys: {key: -> {"key":
    text = re.sub(r"{\s*([a-zA-Z_][a-zA-Z0-9_]*)\s*:", r'{"\1":', text)
    text = re.sub(r",\s*([a-zA-Z_][a-zA-Z0-9_]*)\s*:", r', "\1":', text)

    # Remove trailing commas before } or ]
    text = re.sub(r",\s*}", "}", text)
    text = re.sub(r",\s*]", "]", text)

    # Try to balance brackets if truncated
    open_braces = text.count("{") - text.count("}")
    open_brackets = text.count("[") - text.count("]")

    if open_braces > 0 or open_brackets > 0:
        # Truncated JSON - try to close it
        # Add closing brackets in reverse order of typical nesting
        text = text.rstrip()
        # Remove trailing comma if present
        if text.endswith(","):
            text = text[:-1]
        # Close arrays first, then objects
        text += "]" * open_brackets
        text += "}" * open_braces

    # Validate the repair worked
    try:
        json.loads(text)
        return text
    except json.JSONDecodeError:
        pass

    # Second attempt: more aggressive cleanup
    # Try extracting just the first complete JSON object/array
    depth = 0
    start_idx = -1
    in_string = False
    escape_next = False

    for i, char in enumerate(text):
        if escape_next:
            escape_next = False
            continue
        if char == "\\":
            escape_next = True
            continue
        if char == '"' and not escape_next:
            in_string = not in_string
            continue
        if in_string:
            continue
        if char in "{[":
            if depth == 0:
                start_idx = i
            depth += 1
        elif char in "}]":
            depth -= 1
            if depth == 0 and start_idx >= 0:
                candidate = text[start_idx : i + 1]
                try:
                    json.loads(candidate)
                    return candidate
                except json.JSONDecodeError:
                    pass

    return None


def _classify_soql_error(error_message: str) -> str:
    """Classify a Salesforce error into a repair taxonomy category.

    Returns:
        One of: FIXABLE_FIELD, FIXABLE_RELATIONSHIP, FIXABLE_SYNTAX,
        FIXABLE_AGGREGATION, NOT_FIXABLE, or UNKNOWN
    """
    error_lower = error_message.lower()

    for category, patterns in SOQL_ERROR_TAXONOMY.items():
        for pattern in patterns:
            if pattern.lower() in error_lower:
                return category

    return "UNKNOWN"


def _extract_invalid_fields(error_message: str) -> List[str]:
    """Extract invalid field names from a Salesforce error message.

    Handles various error formats like:
    - "No such column 'FakeField' on entity 'Opportunity'"
    - "INVALID_FIELD: FakeField"
    - "Unknown field: Account.BadField"
    """
    fields: List[str] = []
    error_lower = error_message.lower()

    # Pattern 1: "No such column 'FieldName'"
    import re

    col_pattern = re.compile(r"no such column\s*['\"]?([a-zA-Z0-9_.]+)['\"]?", re.I)
    for match in col_pattern.finditer(error_message):
        fields.append(match.group(1))

    # Pattern 2: "INVALID_FIELD: FieldName" or "Invalid field: FieldName"
    invalid_pattern = re.compile(r"invalid[_ ]field[:\s]+['\"]?([a-zA-Z0-9_.]+)['\"]?", re.I)
    for match in invalid_pattern.finditer(error_message):
        fields.append(match.group(1))

    # Pattern 3: "Unknown field: Object.Field"
    unknown_pattern = re.compile(r"unknown field[:\s]+['\"]?([a-zA-Z0-9_.]+)['\"]?", re.I)
    for match in unknown_pattern.finditer(error_message):
        fields.append(match.group(1))

    # Pattern 4: "Field 'FieldName' does not exist"
    not_exist_pattern = re.compile(r"field\s*['\"]?([a-zA-Z0-9_.]+)['\"]?\s*(?:does not exist|doesn't exist)", re.I)
    for match in not_exist_pattern.finditer(error_message):
        fields.append(match.group(1))

    # Deduplicate while preserving order
    seen = set()
    unique_fields = []
    for f in fields:
        if f.lower() not in seen:
            seen.add(f.lower())
            unique_fields.append(f)

    return unique_fields


def _is_repairable_error(error_message: str) -> bool:
    """Check if a Salesforce error is potentially repairable by the LLM."""
    category = _classify_soql_error(error_message)
    return category.startswith("FIXABLE_")


@dataclass
class SchemaField:
    """Minimal field metadata for schema lookup and validation."""

    name: str
    label: str
    data_type: str
    relationship_name: Optional[str] = None
    reference_to: List[str] = field(default_factory=list)
    help_text: Optional[str] = None
    queryable: bool = True
    accessible: bool = True
    picklist_values: List[str] = field(default_factory=list)
    length: Optional[int] = None

    def tokens(self) -> List[str]:
        basis = f"{self.name} {self.label} {self.help_text or ''}"
        return re.findall(r"[a-z0-9]+", basis.lower())


@dataclass
class SchemaObject:
    """Queryable Salesforce object metadata."""

    name: str
    label: str
    label_plural: str
    fields: Dict[str, SchemaField] = field(default_factory=dict)
    relationships: Dict[str, List[str]] = field(default_factory=dict)

    def all_field_names(self) -> List[str]:
        return sorted(self.fields.keys())


@dataclass
class PromptObject:
    name: str
    label: str
    sample_fields: List[str]
    relationship_hints: List[str]


@dataclass
class SchemaPromptContext:
    """Schema subset passed to the planner."""

    prompt_objects: List[PromptObject]
    tokens: List[str]
    objects_by_name: Dict[str, SchemaObject] = field(default_factory=dict)

    def to_prompt_block(self) -> str:
        lines: List[str] = []
        for obj in self.prompt_objects:
            lines.append(f"{obj.name} ({obj.label})")
            if obj.sample_fields:
                lines.append("  Fields: " + ", ".join(obj.sample_fields))
            if obj.relationship_hints:
                lines.append("  Relationships: " + "; ".join(obj.relationship_hints))
        return "\n".join(lines)

    def to_payload(self) -> Dict[str, Any]:
        return {
            "objects": [
                {
                    "name": obj.name,
                    "label": obj.label,
                    "sampleFields": obj.sample_fields,
                    "relationships": obj.relationship_hints,
                }
                for obj in self.prompt_objects
            ],
            "tokens": self.tokens,
        }


class PlanJoin(BaseModel):
    """Join description emitted by the planner."""

    model_config = ConfigDict(extra="forbid", populate_by_name=True)

    from_object: str = Field(alias="from")
    to_object: str = Field(alias="to")
    via: str
    relationship: Optional[str] = None

    @field_validator("from_object", "to_object", "via")
    @classmethod
    def _non_empty(cls, value: str) -> str:
        if not value or not str(value).strip():
            raise ValueError("Join fields must be non-empty")
        return str(value).strip()


class PlanFilter(BaseModel):
    """Filter predicate emitted by the planner."""

    model_config = ConfigDict(extra="forbid")

    field: str
    op: str = Field(default="=")
    value: Optional[Any] = None
    values: Optional[List[Any]] = None

    @model_validator(mode="before")
    @classmethod
    def _alias_operator(cls, values: Dict[str, Any]) -> Dict[str, Any]:
        if isinstance(values, dict):
            if "operator" in values and "op" not in values:
                values = dict(values)
                values["op"] = values.pop("operator")
        return values

    @field_validator("field")
    @classmethod
    def _ensure_field(cls, value: str) -> str:
        if not value or not str(value).strip():
            raise ValueError("Filter field is required")
        return str(value).strip()

    @field_validator("value", mode="before")
    @classmethod
    def _coerce_value(cls, value: Any) -> Any:
        if isinstance(value, str):
            lowered = value.strip().lower()
            if lowered in {"true", "false"}:
                return lowered == "true"
            if lowered.isdigit():
                try:
                    return int(lowered)
                except ValueError:
                    return value
        return value

    @field_validator("values", mode="before")
    @classmethod
    def _coerce_values(cls, value: Any) -> Any:
        if isinstance(value, list):
            return [cls._coerce_value(item) for item in value]
        return value

    @field_validator("op", mode="before")
    @classmethod
    def _normalize_op(cls, value: Any) -> str:
        op = re.sub(r"\s+", " ", str(value or "=")).strip().upper().replace(" ", "_")
        synonyms = {
            "EQ": "=",
            "EQUALS": "=",
            "EQUAL": "=",
            "==": "=",  # Python/JS style
            "IS": "=",  # Natural language logic
            "NE": "!=",
            "NOT_EQUALS": "!=",
            "NOT_EQUAL": "!=",
            "NOT_EQUAL_TO": "!=",
            "<>": "!=",  # Standard SQL not-equals
            "GT": ">",
            "LT": "<",
            "GTE": ">=",
            "LTE": "<=",
            "GE": ">=",
            "LE": "<=",
            "INCLUDES": "IN",
            "EXCLUDES": "NOT_IN",
        }
        op = synonyms.get(op, op)
        if op in {"TRUE", "FALSE"}:
            return op
        allowed = {
            "=",
            "!=",
            ">",
            "<",
            ">=",
            "<=",
            "IN",
            "NOT_IN",
            "CONTAINS",
            "STARTS_WITH",
            "ENDS_WITH",
            "LIKE",
            "DATE_LITERAL",
            "IS_NULL",
            "IS_NOT_NULL",
            "BETWEEN",
        }
        if op not in allowed:
            raise ValueError(f"Unsupported operator '{op}'")
        return op

    @model_validator(mode="after")
    def _coerce_boolean_literal(self) -> "PlanFilter":
        if isinstance(self.op, str) and self.op in {"TRUE", "FALSE"}:
            bool_value = self.op == "TRUE"
            self.op = "="
            if self.values:
                self.values = [bool_value]
            else:
                self.value = bool_value
        return self


class PlanOrder(BaseModel):
    """Order by clause."""

    model_config = ConfigDict(extra="forbid")

    field: str
    dir: str = Field(default="ASC")

    @field_validator("field")
    @classmethod
    def _ensure_field(cls, value: str) -> str:
        if not value or not str(value).strip():
            raise ValueError("Order field is required")
        return str(value).strip()

    @field_validator("dir", mode="before")
    @classmethod
    def _normalize_dir(cls, value: Any) -> str:
        direction = str(value or "ASC").upper()
        return "DESC" if direction == "DESC" else "ASC"


class SalesforcePlan(BaseModel):
    """Structured query plan DSL used for compilation and validation."""

    model_config = ConfigDict(extra="forbid", populate_by_name=True)

    objects: List[str] = Field(default_factory=list)
    fields: List[str] = Field(default_factory=list)
    joins: List[PlanJoin] = Field(default_factory=list)
    filters: List[PlanFilter] = Field(default_factory=list)
    order_by: List[PlanOrder] = Field(default_factory=list, alias="orderBy")
    group_by: List[str] = Field(default_factory=list, alias="groupBy")
    limit: Optional[int] = None

    @field_validator("objects", "fields", "group_by", mode="before")
    @classmethod
    def _empty_list(cls, value: Any) -> List[str]:
        if value is None:
            return []
        return list(value)


class SuggestedObject(BaseModel):
    """Lightweight object suggestion returned by the researcher stage."""

    model_config = ConfigDict(extra="ignore")

    name: str
    why: Optional[str] = None
    score: Optional[float] = None


class SuggestedField(BaseModel):
    """Field recommendation with rationale."""

    model_config = ConfigDict(extra="ignore", populate_by_name=True)

    obj: str = Field(alias="object")
    field: str
    why: Optional[str] = None
    score: Optional[float] = None
    field_ref: Optional[str] = Field(default=None, alias="fieldRef")
    object_api: Optional[str] = Field(default=None, alias="objectApi")
    report_usage: Optional[int] = Field(default=None, alias="reportUsage")
    fill_rate: Optional[float] = Field(default=None, alias="fillRate")
    badges: List[str] = Field(default_factory=list)

    @field_validator("obj", "field", mode="before")
    @classmethod
    def _strip(cls, value: Any) -> str:
        if value is None:
            return ""
        return str(value).strip()


class SuggestedFilter(BaseModel):
    """Suggested filter snippet (plain text)."""

    model_config = ConfigDict(extra="ignore", populate_by_name=True)

    expr: str = Field(alias="expression")
    why: Optional[str] = None

    @field_validator("expr", mode="before")
    @classmethod
    def _ensure_str(cls, value: Any) -> str:
        return str(value or "").strip()


class ExamplePlan(BaseModel):
    """Example plan concept returned by discovery stage."""

    model_config = ConfigDict(extra="ignore")

    title: str
    summary: Optional[str] = None
    plan: Optional[Dict[str, Any]] = None


class MacroDefinition(BaseModel):
    """Keyboard macro / command palette entry."""

    model_config = ConfigDict(extra="ignore")

    id: str
    title: str
    patch: PlanPatch
    description: Optional[str] = None
    category: Optional[str] = None


class ResearchArtifacts(BaseModel):
    """Aggregated discovery hints from researcher stage."""

    model_config = ConfigDict(extra="ignore")

    clarifying_questions: List[str] = Field(default_factory=list)
    related_objects: List[SuggestedObject] = Field(default_factory=list)
    suggested_fields: List[SuggestedField] = Field(default_factory=list)
    suggested_filters: List[SuggestedFilter] = Field(default_factory=list)
    common_filters: List[str] = Field(default_factory=list)
    example_plans: List[ExamplePlan] = Field(default_factory=list)
    macros: List[MacroDefinition] = Field(default_factory=list)
    ranked_fields: List[Dict[str, Any]] = Field(default_factory=list)
    clarifiers: List[Dict[str, Any]] = Field(default_factory=list)
    index_status: Optional[Dict[str, Any]] = Field(default=None, alias="indexStatus")

    @field_validator(
        "clarifying_questions",
        "common_filters",
        mode="before",
    )
    @classmethod
    def _ensure_list_of_str(cls, value: Any) -> List[str]:
        if value is None:
            return []
        if isinstance(value, list):
            return [str(item).strip() for item in value if str(item).strip()]
        if isinstance(value, str) and value.strip():
            return [value.strip()]
        return []


class PlanPatchOp(BaseModel):
    """Single JSON-Patch style operation."""

    model_config = ConfigDict(extra="forbid")

    op: str
    path: str
    value: Any = None

    @field_validator("op", mode="before")
    @classmethod
    def _normalize_op(cls, value: Any) -> str:
        return str(value or "").lower()


class PlanPatch(BaseModel):
    """Collection of patch operations with optional summaries."""

    model_config = ConfigDict(extra="forbid")

    ops: List[PlanPatchOp] = Field(default_factory=list)
    summary: List[str] = Field(default_factory=list)

    @field_validator("summary", mode="before")
    @classmethod
    def _ensure_summary(cls, value: Any) -> List[str]:
        if value is None:
            return []
        if isinstance(value, list):
            return [str(item) for item in value if str(item).strip()]
        if isinstance(value, str) and value.strip():
            return [value.strip()]
        return []


class PlanAlternative(BaseModel):
    """Critic-proposed alternative modification."""

    model_config = ConfigDict(extra="ignore")

    title: str
    patch: PlanPatch
    description: Optional[str] = None


class PlanClarifierOption(BaseModel):
    """Selectable option for a clarifying question."""

    model_config = ConfigDict(extra="ignore")

    label: str
    value: Optional[str] = None
    patch: Optional[PlanPatch] = None
    hint: Optional[str] = None

    @field_validator("label", mode="before")
    @classmethod
    def _normalize_label(cls, value: Any) -> str:
        return str(value or "").strip()


class PlanClarifier(BaseModel):
    """Clarifying question that can drive structured changes via patches."""

    model_config = ConfigDict(extra="ignore")

    id: str
    question: str
    options: List[PlanClarifierOption] = Field(default_factory=list)
    required: bool = True
    default_option: Optional[str] = Field(default=None, alias="default")
    blocking: bool = False

    @field_validator("question", mode="before")
    @classmethod
    def _normalize_question(cls, value: Any) -> str:
        return str(value or "").strip()

    @field_validator("options", mode="before")
    @classmethod
    def _validate_options(cls, value: Any) -> List[Any]:
        if value is None:
            return []
        if isinstance(value, list):
            return value
        return [value]


class PlannerDraft(BaseModel):
    """Primary planner output before critic review."""

    model_config = ConfigDict(populate_by_name=True)

    plan: SalesforcePlan
    trace: List[str] = Field(default_factory=list)
    clarifiers: List[PlanClarifier] = Field(default_factory=list)

    @field_validator("trace", mode="before")
    @classmethod
    def _ensure_trace(cls, value: Any) -> List[str]:
        if value is None:
            return []
        if isinstance(value, list):
            return [str(item).strip() for item in value if str(item).strip()]
        if isinstance(value, str) and value.strip():
            return [value.strip()]
        return []


class CriticFeedback(BaseModel):
    """Critic stage feedback with risks and alternatives."""

    model_config = ConfigDict(extra="ignore")

    risks: List[str] = Field(default_factory=list)
    alternatives: List[PlanAlternative] = Field(default_factory=list)

    @field_validator("risks", mode="before")
    @classmethod
    def _ensure_risks(cls, value: Any) -> List[str]:
        if value is None:
            return []
        if isinstance(value, list):
            return [str(item).strip() for item in value if str(item).strip()]
        if isinstance(value, str) and value.strip():
            return [value.strip()]
        return []


class PlanPipelineArtifacts(BaseModel):
    """Final payload emitted by multi-agent pipeline."""

    model_config = ConfigDict(extra="ignore")

    plan: SalesforcePlan
    trace: List[str] = Field(default_factory=list)
    clarifiers: List[PlanClarifier] = Field(default_factory=list)
    risks: List[str] = Field(default_factory=list)
    alternatives: List[PlanAlternative] = Field(default_factory=list)
    research: Optional[ResearchArtifacts] = None


# ==================== JSON Schemas for Gemini response_schema ====================
# These schemas are passed to Vertex AI to enforce structured JSON output.
# Generated from Pydantic models for type safety and consistency.

def _sanitize_vertex_schema(schema: Dict[str, Any]) -> Dict[str, Any]:
    """Strip null-only schema branches that Vertex rejects."""

    def _is_null_type(value: Any) -> bool:
        if value is None:
            return True
        if isinstance(value, str):
            return value.lower() == "null"
        if isinstance(value, list):
            return all(_is_null_type(item) for item in value)
        return False

    def _strip(node: Any) -> Any:
        if isinstance(node, list):
            return [_strip(item) for item in node if item is not None]
        if not isinstance(node, dict):
            return node

        cleaned = {key: _strip(value) for key, value in node.items()}
        node_type = cleaned.get("type")
        if _is_null_type(node_type):
            cleaned.pop("type", None)
        elif isinstance(node_type, list):
            filtered = [item for item in node_type if not _is_null_type(item)]
            if len(filtered) == 1:
                cleaned["type"] = filtered[0]
            else:
                cleaned["type"] = filtered or node_type

        any_of = cleaned.get("anyOf")
        if isinstance(any_of, list):
            filtered_any = []
            for item in any_of:
                if isinstance(item, dict) and _is_null_type(item.get("type")):
                    continue
                filtered_any.append(item)
            if len(filtered_any) == 1 and isinstance(filtered_any[0], dict):
                cleaned.pop("anyOf", None)
                for key, value in filtered_any[0].items():
                    cleaned.setdefault(key, value)
            else:
                cleaned["anyOf"] = filtered_any

        return cleaned

    return _strip(schema)


def _generate_schema_safe(model_cls: type) -> Optional[Dict[str, Any]]:
    """
    Generate JSON schema from Pydantic model, returning None on failure.

    This allows graceful degradation if schema generation fails for any reason
    (e.g., unsupported type, circular reference).
    """
    try:
        schema = model_cls.model_json_schema()
        return _sanitize_vertex_schema(schema)
    except Exception as exc:
        log.warning("Failed to generate schema for %s: %s", model_cls.__name__, exc)
        return None


# Pre-generate schemas at module load time
RESEARCHER_SCHEMA: Optional[Dict[str, Any]] = _generate_schema_safe(ResearchArtifacts)
PLANNER_SCHEMA: Optional[Dict[str, Any]] = _generate_schema_safe(PlannerDraft)
CRITIC_SCHEMA: Optional[Dict[str, Any]] = _generate_schema_safe(CriticFeedback)

# Log schema availability
if RESEARCHER_SCHEMA:
    log.debug("ResearchArtifacts schema available for Gemini response_schema")
if PLANNER_SCHEMA:
    log.debug("PlannerDraft schema available for Gemini response_schema")
if CRITIC_SCHEMA:
    log.debug("CriticFeedback schema available for Gemini response_schema")


def _tokenize(text: str) -> List[str]:
    return [
        tok for tok in re.findall(r"[a-z0-9]+", (text or "").lower()) if len(tok) > 2
    ]


class SalesforceSchemaIndex:
    """Caches per-tenant Salesforce schema metadata for AI planning."""

    def __init__(self) -> None:
        self._global_cache = HybridCache(
            namespace="sf_schema_global", lru_size=32, ttl=AI_REPORTS_SCHEMA_CACHE_TTL
        )
        self._object_cache = HybridCache(
            namespace="sf_schema_object", lru_size=128, ttl=AI_REPORTS_SCHEMA_CACHE_TTL
        )

    async def _fetch_global(
        self, account_id: str, gateway: SalesforceGateway
    ) -> Dict[str, Any]:
        cache_key = f"global:{account_id}"
        cached = await self._global_cache.get(cache_key)
        if cached:
            return cached
        payload = await gateway.describe_global()
        data = {
            "sobjects": [
                {
                    "name": obj.get("name"),
                    "label": obj.get("label"),
                    "labelPlural": obj.get("labelPlural") or obj.get("label"),
                    "queryable": obj.get("queryable", False),
                    "searchable": obj.get("searchable", False),
                    "custom": obj.get("custom", False),
                    "keyPrefix": obj.get("keyPrefix"),
                }
                for obj in payload.get("sobjects", [])
            ]
        }
        await self._global_cache.set(cache_key, data)
        return data

    async def _fetch_object(
        self, account_id: str, object_name: str, gateway: SalesforceGateway
    ) -> SchemaObject:
        normalized = object_name.strip()
        cache_key = f"obj:{account_id}:{normalized.lower()}"
        cached = await self._object_cache.get(cache_key)
        if cached:
            return cached  # type: ignore[return-value]

        raw = await gateway.describe_object(normalized)
        schema_object = SchemaObject(
            name=normalized,
            label=raw.get("label", normalized),
            label_plural=raw.get("labelPlural", raw.get("label", normalized) + "s"),
        )
        for field in raw.get("fields", []):
            name = field.get("name")
            if not name:
                continue
            schema_field = SchemaField(
                name=name,
                label=field.get("label", name),
                data_type=field.get("type", ""),
                relationship_name=field.get("relationshipName"),
                reference_to=list(field.get("referenceTo") or []),
                help_text=field.get("inlineHelpText"),
                queryable=field.get("queryable", True),
                accessible=field.get(
                    "accessible", field.get("permissions", {}).get("readable", True)
                ),
                picklist_values=[
                    entry.get("value")
                    for entry in field.get("picklistValues", [])
                    if isinstance(entry, dict)
                    and entry.get("active", True)
                    and entry.get("value")
                ],
                length=field.get("length"),
            )
            schema_object.fields[name.lower()] = schema_field
            rel = field.get("relationshipName")
            if rel:
                schema_object.relationships[rel.lower()] = [
                    ref for ref in field.get("referenceTo", []) if ref
                ]
        await self._object_cache.set(cache_key, schema_object)
        return schema_object

    async def build_prompt_context(
        self,
        *,
        account_id: str,
        gateway: SalesforceGateway,
        prompt: str,
        max_objects: int,
        max_fields: int,
    ) -> SchemaPromptContext:
        tokens = _tokenize(prompt)
        global_schema = await self._fetch_global(account_id, gateway)
        objects = [
            obj for obj in global_schema.get("sobjects", []) if obj.get("queryable")
        ]

        allow_objects = {obj.lower() for obj in AI_REPORTS_ALLOW_OBJECTS or []}
        deny_objects = {obj.lower() for obj in AI_REPORTS_DENY_OBJECTS or []}

        if allow_objects:
            objects = [
                obj
                for obj in objects
                if str(obj.get("name", "")).lower() in allow_objects
            ]
        if deny_objects:
            objects = [
                obj
                for obj in objects
                if str(obj.get("name", "")).lower() not in deny_objects
            ]

        if not objects:
            return SchemaPromptContext(
                prompt_objects=[], tokens=tokens, objects_by_name={}
            )

        def score_object(entry: Dict[str, Any]) -> float:
            base = 0.5 if not entry.get("custom") else 0.3
            label_text = f"{entry.get('label', '')} {entry.get('name', '')}".lower()
            for token in tokens:
                if token in (entry.get("name") or "").lower():
                    base += 2.0
                if token in label_text:
                    base += 3.0
            return base

        ranked = sorted(objects, key=score_object, reverse=True)
        if not tokens:
            priority_names = ["Opportunity", "Lead", "Contact", "Account", "Case"]
            ranked = sorted(
                objects,
                key=lambda obj: (
                    0 if obj.get("name") in priority_names else 1,
                    obj.get("label", ""),
                ),
            )
        # Ensure critical objects (e.g. Opportunity for CloseDate filters) stay in context
        required_objects = ["Opportunity"]
        for target in reversed(required_objects):
            for idx, entry in enumerate(ranked):
                if entry.get("name") == target:
                    ranked.insert(0, ranked.pop(idx))
                    break

        prompt_objects: List[PromptObject] = []
        objects_by_name: Dict[str, SchemaObject] = {}

        for entry in ranked[:max_objects]:
            name = entry.get("name")
            if not name:
                continue
            schema_object = await self._fetch_object(account_id, name, gateway)
            field_scores: List[Tuple[float, SchemaField]] = []
            for field in schema_object.fields.values():
                score = 0.2
                if field.name.lower() in {"id", "name"}:
                    score += 4.0
                for token in tokens:
                    if token in field.name.lower():
                        score += 1.5
                    if token in field.label.lower():
                        score += 2.0
                field_scores.append((score, field))

            selected_fields = [
                f"{schema_object.name}.{field.name}"
                for _, field in sorted(
                    field_scores, key=lambda item: item[0], reverse=True
                )[:max_fields]
            ]

            relationship_hints: List[str] = []
            for rel_name, targets in list(schema_object.relationships.items())[:4]:
                rel_target = ", ".join(targets[:2])
                relationship_hints.append(f"{rel_name} -> {rel_target}")

            prompt_objects.append(
                PromptObject(
                    name=schema_object.name,
                    label=schema_object.label,
                    sample_fields=selected_fields,
                    relationship_hints=relationship_hints,
                )
            )

            objects_by_name[schema_object.name.lower()] = schema_object

        return SchemaPromptContext(
            prompt_objects=prompt_objects,
            tokens=tokens,
            objects_by_name=objects_by_name,
        )


# Pattern to detect potentially suspicious SOQL fragments in user input
_SOQL_SUSPICIOUS_PATTERN = re.compile(
    r"(?i)(\bSELECT\b|\bFROM\b|\bWHERE\b|\bLIMIT\b|\bOFFSET\b|--|/\*|\*/)"
)


_SOQL_DATE_LITERAL_PATTERN = re.compile(
    r"^(?:"
    r"TODAY|YESTERDAY|TOMORROW|"
    r"(?:LAST|THIS|NEXT)_(?:WEEK|MONTH|QUARTER|YEAR|"
    r"FISCAL_QUARTER|FISCAL_YEAR|90_DAYS)|"
    r"(?:LAST|NEXT)_N_(?:DAYS|WEEKS|MONTHS|QUARTERS|YEARS|"
    r"FISCAL_QUARTERS|FISCAL_YEARS):\d+"
    r")$",
    re.IGNORECASE,
)

_SOQL_STRICT_DATE_VALUE_PATTERN = re.compile(
    r"^\d{4}-\d{2}-\d{2}(T\d{2}:\d{2}:\d{2}(\.\d+)?(Z|[+-]\d{2}:\d{2})?)?$"
)

_SOQL_IDENTIFIER_PATTERN = re.compile(r"^[A-Za-z_][A-Za-z0-9_.]*$")
_SOQL_AGGREGATE_FIELD_PATTERN = re.compile(
    r"(?i)^(?:sum|avg|count|min|max|count_distinct)\s*\(\s*([A-Za-z_][A-Za-z0-9_.]*|\*)\s*\)$"
)


def _escape_soql_literal(raw: str) -> str:
    escaped = str(raw)
    escaped = escaped.replace("\\", "\\\\")
    escaped = escaped.replace("'", "\\'")
    escaped = escaped.replace("\n", "\\n")
    escaped = escaped.replace("\r", "\\r")
    escaped = escaped.replace("\t", "\\t")
    escaped = escaped.replace("\0", "")
    return escaped


def _validate_soql_identifier(name: str, label: str) -> str:
    text = str(name or "").strip()
    if not _SOQL_IDENTIFIER_PATTERN.fullmatch(text):
        raise PlanValidationError(f"Invalid {label} identifier '{text}'")
    return text


def _validate_soql_field_expr(expr: str, label: str = "field") -> str:
    text = str(expr or "").strip()
    if not text:
        raise PlanValidationError(f"Invalid {label} identifier '{text}'")
    agg_match = _SOQL_AGGREGATE_FIELD_PATTERN.fullmatch(text)
    if agg_match:
        inner = agg_match.group(1)
        if inner != "*":
            _validate_soql_identifier(inner, label)
        return text
    _validate_soql_identifier(text, label)
    return text


def _format_value(value: Any) -> str:
    """Format a Python value as a SOQL literal with comprehensive escaping.

    Handles None, booleans, numbers, and strings with proper escaping
    for SOQL injection prevention. SOQL date literals (e.g., LAST_N_DAYS:14)
    are returned unquoted.
    """
    if value is None:
        return "null"
    if isinstance(value, bool):
        return "true" if value else "false"
    if isinstance(value, (int, float)):
        return str(value)

    raw = str(value)

    # SOQL date literals must NOT be quoted
    if _SOQL_DATE_LITERAL_PATTERN.match(raw):
        return raw

    # Log suspicious patterns (but don't block - legitimate data may contain these)
    if _SOQL_SUSPICIOUS_PATTERN.search(raw):
        log.warning(
            "SOQL filter value contains suspicious pattern: %s",
            raw[:100] if len(raw) > 100 else raw,
        )

    return f"'{_escape_soql_literal(raw)}'"


def _hash_plan_dict(payload: Dict[str, Any]) -> str:
    data = json.dumps(payload, sort_keys=True, separators=(",", ":")).encode("utf-8")
    return hashlib.sha256(data).hexdigest()


def _collect_allowlist(
    plan: SalesforcePlan, schema_context: Optional[SchemaPromptContext] = None
) -> Dict[str, Any]:
    fields_by_obj: Dict[str, set] = {}
    objects: set[str] = set()
    accessible_by_obj: Optional[Dict[str, set[str]]] = None
    if schema_context is not None:
        accessible_by_obj = {}
        for schema_obj in schema_context.objects_by_name.values():
            accessible_by_obj[schema_obj.name.lower()] = {
                field.name.lower()
                for field in schema_obj.fields.values()
                if field.accessible and field.queryable
            }

    def add_object(object_name: str) -> None:
        key = str(object_name or "").strip().lower()
        if not key:
            return
        if accessible_by_obj is not None and key not in accessible_by_obj:
            return
        objects.add(key)
        fields_by_obj.setdefault(key, set())

    def add_field(object_name: str, field_name: str) -> None:
        key = str(object_name or "").strip().lower()
        field_key = str(field_name or "").strip().lower()
        if not key or not field_key:
            return
        if accessible_by_obj is not None:
            allowed = accessible_by_obj.get(key)
            if allowed is None or field_key not in allowed:
                return
        add_object(key)
        fields_by_obj.setdefault(key, set()).add(field_key)

    primary = plan.objects[0] if plan.objects else "Lead"
    add_object(primary)
    for obj_name in plan.objects or []:
        add_object(obj_name)
    for join in plan.joins or []:
        add_object(join.from_object)
        add_object(join.to_object)
        if join.relationship:
            add_object(join.relationship)

    if not plan.fields:
        add_field(primary, "id")
    for field in plan.fields:
        normalized = _unwrap_field_reference(field)
        if not normalized:
            continue
        target_obj = _resolve_object_from_path(normalized, primary)
        add_field(target_obj, normalized.split(".")[-1])
    for filt in plan.filters:
        normalized = _unwrap_field_reference(filt.field)
        if not normalized:
            continue
        target_obj = _resolve_object_from_path(normalized, primary)
        add_field(target_obj, normalized.split(".")[-1])
    for order in plan.order_by:
        normalized = _unwrap_field_reference(order.field)
        if not normalized:
            continue
        target_obj = _resolve_object_from_path(normalized, primary)
        add_field(target_obj, normalized.split(".")[-1])
    for group_field in plan.group_by:
        normalized = _unwrap_field_reference(group_field)
        if not normalized:
            continue
        target_obj = _resolve_object_from_path(normalized, primary)
        add_field(target_obj, normalized.split(".")[-1])
    return {
        "objects": sorted(objects),
        "fields": {obj: sorted(fields) for obj, fields in fields_by_obj.items()},
    }


def _resolve_object_from_path(field_path: str, fallback_object: str) -> str:
    if not field_path:
        return fallback_object
    parts = field_path.split(".")
    if len(parts) == 1:
        return fallback_object
    # If the first part looks like an object name (starts with uppercase), treat it as such
    first = parts[0]
    if first[:1].isupper():
        return first
    return fallback_object


def _unwrap_field_reference(field_expr: str) -> str:
    text = str(field_expr or "").strip()
    if not text:
        return ""
    agg_match = _SOQL_AGGREGATE_FIELD_PATTERN.fullmatch(text)
    if agg_match:
        inner = agg_match.group(1).strip()
        if inner == "*":
            return ""
        return inner
    return text


def _pointer_unescape(token: str) -> str:
    return token.replace("~1", "/").replace("~0", "~")


def _pointer_tokens(path: str) -> List[str]:
    if path in ("", "/"):
        return []
    if not path.startswith("/"):
        raise PlanValidationError(f"Patch path must start with '/': {path}")
    parts = path.lstrip("/").split("/")
    return [
        _pointer_unescape(part) for part in parts if part != "" or path.endswith("/")
    ]


def _ensure_list_length(target: List[Any], index: int) -> None:
    while len(target) <= index:
        target.append({})


def _resolve_parent(
    doc: Union[Dict[str, Any], List[Any]],
    tokens: List[str],
    *,
    create: bool,
) -> Tuple[Union[Dict[str, Any], List[Any]], str]:
    if not tokens:
        return doc, ""
    node: Union[Dict[str, Any], List[Any]] = doc
    for i, token in enumerate(tokens[:-1]):
        if isinstance(node, dict):
            if token not in node:
                if not create:
                    raise PlanValidationError(
                        f"Missing path element '{token}' in patch"
                    )
                node[token] = {}
            node = node[token]
        elif isinstance(node, list):
            if token == "-":
                raise PlanValidationError(
                    "'-' token only allowed at final position in patch path"
                )
            try:
                idx = int(token)
            except ValueError as exc:
                raise PlanValidationError(
                    f"Invalid list index '{token}' in patch path"
                ) from exc
            if idx < 0:
                raise PlanValidationError(
                    f"Negative list index '{idx}' not supported in patches"
                )
            if idx >= len(node):
                if not create:
                    raise PlanValidationError(
                        f"List index {idx} out of range for path segment '{token}'"
                    )
                _ensure_list_length(node, idx)
            node = node[idx]
        else:
            raise PlanValidationError(
                "Cannot traverse non-container type in patch path"
            )
    return node, tokens[-1]


def _apply_patch_operations(
    plan_payload: Dict[str, Any],
    patch: PlanPatch,
) -> Dict[str, Any]:
    updated = deepcopy(plan_payload)
    for op in patch.ops:
        tokens = _pointer_tokens(op.path)
        if not tokens:
            parent = updated
            key = ""
        else:
            parent, key = _resolve_parent(updated, tokens, create=op.op == "add")

        if op.op == "add":
            if not tokens:
                if not isinstance(op.value, dict):
                    raise PlanValidationError(
                        "Root add operation must provide object value"
                    )
                updated = deepcopy(op.value)
                continue
            if isinstance(parent, list):
                if key == "-":
                    parent.append(deepcopy(op.value))
                else:
                    try:
                        idx = int(key)
                    except ValueError as exc:
                        raise PlanValidationError(
                            f"Invalid list index '{key}' for add"
                        ) from exc
                    if idx < 0 or idx > len(parent):
                        raise PlanValidationError(
                            f"Add index {idx} out of range for list of size {len(parent)}"
                        )
                    parent.insert(idx, deepcopy(op.value))
            elif isinstance(parent, dict):
                parent[key] = deepcopy(op.value)
            else:
                raise PlanValidationError("Add operation target must be dict or list")
        elif op.op == "replace":
            target = parent if tokens else updated
            if isinstance(target, list):
                if key == "":
                    raise PlanValidationError(
                        "Replace root list without path is not supported"
                    )
                try:
                    idx = int(key)
                except ValueError as exc:
                    raise PlanValidationError(
                        f"Invalid list index '{key}' for replace"
                    ) from exc
                if idx < 0 or idx >= len(target):
                    raise PlanValidationError(
                        f"Replace index {idx} out of range for list"
                    )
                target[idx] = deepcopy(op.value)
            elif isinstance(target, dict):
                if key not in target:
                    raise PlanValidationError(f"Cannot replace missing key '{key}'")
                target[key] = deepcopy(op.value)
            else:
                raise PlanValidationError(
                    "Replace operation target must be dict or list"
                )
        elif op.op == "remove":
            target = parent if tokens else updated
            if isinstance(target, list):
                if key == "":
                    raise PlanValidationError(
                        "Remove root list without path is not supported"
                    )
                try:
                    idx = int(key)
                except ValueError as exc:
                    raise PlanValidationError(
                        f"Invalid list index '{key}' for remove"
                    ) from exc
                if idx < 0 or idx >= len(target):
                    raise PlanValidationError(
                        f"Remove index {idx} out of range for list"
                    )
                target.pop(idx)
            elif isinstance(target, dict):
                if key not in target:
                    raise PlanValidationError(f"Cannot remove missing key '{key}'")
                target.pop(key)
            else:
                raise PlanValidationError(
                    "Remove operation target must be dict or list"
                )
        else:
            raise PlanValidationError(f"Unsupported patch operation '{op.op}'")
    return updated


def _ensure_patch_summary(patch: PlanPatch) -> PlanPatch:
    if patch.summary:
        return patch
    summaries: List[str] = []
    for op in patch.ops:
        if op.op == "add" and op.path.startswith("/fields"):
            summaries.append(f"+ fields: {op.value}")
        elif op.op == "add" and op.path.startswith("/filters"):
            summaries.append(f"+ filters: {op.value}")
        elif op.op == "replace" and op.path.startswith("/limit"):
            summaries.append(f"* limit -> {op.value}")
        elif op.op == "remove" and op.path.startswith("/fields"):
            summaries.append(f"- fields[{op.path.split('/')[-1]}]")
    if not summaries:
        summaries = [f"{op.op.upper()} {op.path}" for op in patch.ops]
    patch.summary = summaries
    return patch


DEFAULT_MACROS: List[MacroDefinition] = [
    MacroDefinition(
        id="add-stage-amount",
        title="Add Stage & Amount",
        description="Add StageName and Amount fields to the selection.",
        patch=PlanPatch(
            ops=[
                PlanPatchOp(op="add", path="/fields/-", value="Opportunity.StageName"),
                PlanPatchOp(op="add", path="/fields/-", value="Opportunity.Amount"),
            ]
        ),
        category="Fields",
    ),
    MacroDefinition(
        id="set-date-this-quarter",
        title="Set date: THIS_QUARTER",
        description="Restrict CloseDate to the current quarter.",
        patch=PlanPatch(
            ops=[
                PlanPatchOp(
                    op="add",
                    path="/filters/-",
                    value={
                        "field": "Opportunity.CloseDate",
                        "op": "DATE_LITERAL",
                        "value": "THIS_QUARTER",
                    },
                ),
            ]
        ),
        category="Filters",
    ),
    MacroDefinition(
        id="group-by-stage",
        title="Group by Stage",
        description="Group the results by StageName.",
        patch=PlanPatch(
            ops=[
                PlanPatchOp(op="add", path="/groupBy/-", value="Opportunity.StageName"),
            ]
        ),
        category="Aggregation",
    ),
    MacroDefinition(
        id="order-by-close-desc",
        title="Sort by CloseDate desc",
        description="Order by CloseDate in descending order.",
        patch=PlanPatch(
            ops=[
                PlanPatchOp(
                    op="add",
                    path="/orderBy/-",
                    value={"field": "Opportunity.CloseDate", "dir": "DESC"},
                ),
            ]
        ),
        category="Sorting",
    ),
    MacroDefinition(
        id="show-primary-contact",
        title="Show primary contact (OCR)",
        description="Include primary contact email via OpportunityContactRole.",
        patch=PlanPatch(
            ops=[
                PlanPatchOp(
                    op="add",
                    path="/joins/-",
                    value={
                        "from": "Opportunity",
                        "to": "OpportunityContactRole",
                        "via": "Id",
                        "relationship": "OpportunityContactRoles",
                    },
                ),
                PlanPatchOp(
                    op="add",
                    path="/fields/-",
                    value="OpportunityContactRoles.Contact.Email",
                ),
                PlanPatchOp(
                    op="add",
                    path="/filters/-",
                    value={
                        "field": "OpportunityContactRoles.IsPrimary",
                        "op": "=",
                        "value": True,
                    },
                ),
            ]
        ),
        category="Relationships",
    ),
]


class SalesforceAIReportService:
    """Coordinates schema context, LLM planning, validation, compilation, and preview.

    Uses Vertex AI exclusively. Ollama support has been removed.
    """

    def __init__(self, schema_index: Optional[SalesforceSchemaIndex] = None) -> None:
        self.schema_index = schema_index or SalesforceSchemaIndex()
        self._planner_clients: Dict[str, VertexAIProvider] = {}
        self._validator_client: Optional[VertexAIProvider] = None
        self._macro_registry: List[MacroDefinition] = [
            macro.model_copy(deep=True) for macro in DEFAULT_MACROS
        ]
        self._embedding_model = None
        self._embeddings_disabled = False
        self._embeddings_disabled_logged = False
        # Context caching: Maps schema content hash -> Vertex cache name
        self._active_caches: Dict[str, str] = {}
        self._cache_lock = threading.Lock()  # Thread-safe cache access
        self._client_lock = threading.Lock()  # Thread-safe client initialization
        self._research_cache = HybridCache(
            namespace="ai_reports_research",
            lru_size=256,
            ttl=AI_REPORTS_RESEARCH_CACHE_TTL,
        )
        self._in_flight_research: Dict[str, asyncio.Task] = {}
        self._research_lock = asyncio.Lock()

    def _get_client(self, model_name: str) -> VertexAIProvider:
        """Get or create a Vertex AI client for the given model.

        Raises:
            PlannerUnavailable: If Vertex AI initialization fails.
        """
        # Check without lock first (fast path for already-initialized clients)
        if model_name in self._planner_clients:
            return self._planner_clients[model_name]

        # Thread-safe initialization for new clients
        with self._client_lock:
            # Double-check after acquiring lock
            if model_name not in self._planner_clients:
                project = (
                    VERTEX_PROJECT
                    or os.getenv("GOOGLE_CLOUD_PROJECT")
                    or os.getenv("GCP_PROJECT")
                )
                location = VERTEX_LOCATION or os.getenv("VERTEX_LOCATION")
                try:
                    self._planner_clients[model_name] = VertexAIProvider(
                        model=model_name,
                        project=project,
                        location=location,
                    )
                    log.info("Using Vertex AI Gemini model: %s", model_name)
                except Exception as exc:
                    log.error(
                        "Vertex AI init failed for %s (%s). "
                        "Ensure VERTEX_PROJECT is set and credentials are configured.",
                        model_name,
                        exc,
                    )
                    raise PlannerUnavailable(
                        f"Vertex AI initialization failed for {model_name}: {exc}. "
                        "Ollama fallback has been removed - configure Vertex AI properly."
                    ) from exc
        return self._planner_clients[model_name]

    def _get_validator(self) -> VertexAIProvider:
        """Get or create a Vertex AI validator client.

        Raises:
            PlannerUnavailable: If Vertex AI initialization fails.
        """
        if not self._validator_client:
            model_name = AI_REPORTS_VALIDATOR_MODEL or AI_REPORTS_PLANNER_MODEL
            project = (
                VERTEX_PROJECT
                or os.getenv("GOOGLE_CLOUD_PROJECT")
                or os.getenv("GCP_PROJECT")
            )
            location = VERTEX_LOCATION or os.getenv("VERTEX_LOCATION")
            try:
                self._validator_client = VertexAIProvider(
                    model=model_name,
                    project=project,
                    location=location,
                )
            except Exception as exc:
                log.error(
                    "Vertex AI validator init failed for %s (%s). "
                    "Ensure VERTEX_PROJECT is set and credentials are configured.",
                    model_name,
                    exc,
                )
                raise PlannerUnavailable(
                    f"Vertex AI validator initialization failed for {model_name}: {exc}. "
                    "Ollama fallback has been removed - configure Vertex AI properly."
                ) from exc
        return self._validator_client

    def _get_or_create_cache(self, schema_context: SchemaPromptContext, client: VertexAIProvider) -> str:
        """
        Get or create a Vertex AI cache for the given schema context.

        Args:
            schema_context: The schema context to cache
            client: The VertexAIProvider client

        Returns:
            Cache name (resource path) or empty string if caching unavailable
        """
        # Generate schema content for caching
        schema_text = schema_context.to_prompt_block()

        # Check if content is large enough to cache (Vertex requires 2048+ tokens)
        # However, token count includes punctuation, field names, formatting, etc.
        # A 300-word schema with field details can easily exceed 2048 tokens
        word_count = len(schema_text.split())
        char_count = len(schema_text)
        approx_tokens = max(word_count, int(char_count / 4))

        # Log the actual schema size for debugging
        log.info(
            "Schema context size: %s words, %s chars (~%s tokens)",
            word_count,
            char_count,
            approx_tokens,
        )

        min_cache_tokens = 2048
        # Vertex requires >=2048 tokens to accept cached content.
        if approx_tokens < min_cache_tokens:
            log.debug(
                "Schema too small to cache (~%s tokens, need %s+). Skipping cache.",
                approx_tokens,
                min_cache_tokens,
            )
            return ""

        # Use MD5 hash as cache key
        cache_key = hashlib.md5(schema_text.encode()).hexdigest()

        # Thread-safe cache access to prevent race conditions
        with self._cache_lock:
            # Return existing cache if available
            if cache_key in self._active_caches:
                log.debug(f"Using existing cache for key: {cache_key[:8]}...")
                return self._active_caches[cache_key]

            # Create new cache
            try:
                cache_name = client.create_cache(
                    content=schema_text,
                    ttl_minutes=60,  # 60-minute TTL
                    display_name=f"sf_schema_{cache_key[:8]}",
                )
                if cache_name:
                    self._active_caches[cache_key] = cache_name
                    log.info(f"Created new cache: {cache_key[:8]}... -> {cache_name} ({word_count} words)")
                return cache_name
            except Exception as exc:
                log.warning(f"Failed to create cache for schema: {exc}")
                return ""

    def list_macros(self) -> List[MacroDefinition]:
        clones: List[MacroDefinition] = [
            macro.model_copy(deep=True) for macro in self._macro_registry
        ]
        for macro in clones:
            macro.patch = _ensure_patch_summary(macro.patch)
        return clones

    @staticmethod
    def tier_row_cap(tier: Optional[str]) -> int:
        if not tier:
            return AI_REPORTS_ROW_CAPS.get("FREE", 5000)
        return AI_REPORTS_ROW_CAPS.get(
            tier.upper(), AI_REPORTS_ROW_CAPS.get("FREE", 5000)
        )

    @staticmethod
    def select_model_for_tier(tier: Optional[str]) -> str:
        if not tier:
            return AI_REPORTS_PLANNER_MODEL
        normalized = tier.upper()
        if normalized in {"PRO", "TEAM"}:
            return AI_REPORTS_PRO_MODEL or AI_REPORTS_PLANNER_MODEL
        if normalized in {"ENTERPRISE", "ENT"}:
            return (
                AI_REPORTS_ENTERPRISE_MODEL
                or AI_REPORTS_PRO_MODEL
                or AI_REPORTS_PLANNER_MODEL
            )
        return AI_REPORTS_PLANNER_MODEL

    async def _build_context_from_index(
        self,
        *,
        org_id: str,
        repo: SalesforceSchemaIndexRepo,
        prompt: str,
    ) -> Optional[SchemaPromptContext]:
        objects = await repo.list_objects(org_id, limit=AI_REPORTS_MAX_OBJECTS * 2)
        if not objects:
            return None

        tokens = _tokenize(prompt)
        prompt_objects: List[PromptObject] = []
        objects_by_name: Dict[str, SchemaObject] = {}

        prioritized_objects: List[Any] = []
        other_objects: List[Any] = []
        for obj in objects:
            if (obj.object_api or "").lower() == "opportunity":
                prioritized_objects.append(obj)
            else:
                other_objects.append(obj)
        final_selection = prioritized_objects + other_objects

        for obj_row in final_selection[:AI_REPORTS_MAX_OBJECTS]:
            fields = await repo.list_fields(org_id, obj_row.object_api)
            if not fields:
                continue
            schema_obj = SchemaObject(
                name=obj_row.object_api,
                label=obj_row.label or obj_row.object_api,
                label_plural=obj_row.label_plural
                or obj_row.label
                or obj_row.object_api,
            )
            field_scores: List[Tuple[float, SchemaField]] = []
            for field_row in fields:
                field_accessible = bool(getattr(field_row, "accessible", True))
                if not field_accessible:
                    continue
                schema_field = SchemaField(
                    name=field_row.field_api,
                    label=field_row.label or field_row.field_api,
                    data_type=field_row.data_type,
                    relationship_name=field_row.relationship_api,
                    reference_to=[
                        ref for ref in (field_row.reference_to or "").split(",") if ref
                    ],
                    help_text=field_row.help_text,
                    queryable=field_row.queryable,
                    accessible=field_accessible,
                    picklist_values=json.loads(field_row.picklist_json or "[]"),
                    length=field_row.length,
                )
                schema_obj.fields[schema_field.name.lower()] = schema_field
                if field_row.relationship_api and field_row.reference_to:
                    schema_obj.relationships[field_row.relationship_api.lower()] = [
                        ref.strip()
                        for ref in field_row.reference_to.split(",")
                        if ref.strip()
                    ]

                score = 0.2
                field_name_lower = schema_field.name.lower()
                label_lower = (schema_field.label or "").lower()

                critical_fields = {
                    "id",
                    "name",
                    "email",
                    "subject",
                    "amount",
                    "stagename",
                    "closedate",
                    "createddate",
                    "activitydate",
                    "status",
                    "description",
                    "type",
                    "priority",
                    "isclosed",
                    "iswon",
                }
                if field_name_lower in critical_fields:
                    score += 5.0
                if not field_row.is_custom:
                    score += 0.5
                if field_row.data_type in {"currency", "double", "percent"}:
                    score += 0.5
                if field_name_lower in {"amount", "expectedrevenue", "annualrevenue"}:
                    score += 2.0
                for token in tokens:
                    if token in field_name_lower:
                        score += 1.2
                    if token in label_lower:
                        score += 1.8
                field_scores.append((score, schema_field))

            if not field_scores:
                continue

            selected_fields = [
                f"{schema_obj.name}.{field.name}"
                for _, field in sorted(
                    field_scores, key=lambda item: item[0], reverse=True
                )[: min(AI_REPORTS_MAX_FIELDS, 20)]
            ]
            relationship_hints = []
            for rel_name, targets in list(schema_obj.relationships.items())[:4]:
                relationship_hints.append(f"{rel_name} -> {', '.join(targets[:2])}")

            prompt_objects.append(
                PromptObject(
                    name=schema_obj.name,
                    label=schema_obj.label,
                    sample_fields=selected_fields,
                    relationship_hints=relationship_hints,
                )
            )
            objects_by_name[schema_obj.name.lower()] = schema_obj

        if not prompt_objects:
            return None

        return SchemaPromptContext(
            prompt_objects=prompt_objects,
            tokens=tokens,
            objects_by_name=objects_by_name,
        )

    def _maybe_disable_embeddings(self, exc: Exception) -> bool:
        message = str(exc).lower()
        quota_hint = "quota project" in message or "aiplatform.googleapis.com" in message
        if not quota_hint:
            return False
        if not self._embeddings_disabled:
            self._embeddings_disabled = True
        if not self._embeddings_disabled_logged:
            log.warning(
                "Vertex AI embeddings disabled until a quota project is set. "
                "Run: gcloud auth application-default set-quota-project <project-id>"
            )
            self._embeddings_disabled_logged = True
        return True

    def _ensure_embedding_model(self):
        if TextEmbeddingModel is None or vertexai is None:
            raise RuntimeError("Vertex AI embeddings unavailable")
        if self._embeddings_disabled:
            return None
        if self._embedding_model is None:
            project = (
                VERTEX_PROJECT
                or os.getenv("VERTEX_PROJECT")
                or os.getenv("GOOGLE_CLOUD_PROJECT")
                or os.getenv("GCP_PROJECT")
            )
            if not project:
                raise RuntimeError("VERTEX_PROJECT is not configured")
            # Use separate location for embeddings (not available in 'global')
            location = os.getenv("VERTEX_EMBEDDING_LOCATION") or os.getenv("VERTEX_LOCATION", "us-central1")
            model_name = os.getenv("VERTEX_EMBEDDING_MODEL", "text-embedding-004")
            try:
                vertexai.init(project=project, location=location)
                self._embedding_model = TextEmbeddingModel.from_pretrained(model_name)
            except Exception as exc:  # noqa: BLE001
                if self._maybe_disable_embeddings(exc):
                    return None
                raise
        return self._embedding_model

    def _schedule_schema_refresh(self, account_id: str) -> None:
        try:
            asyncio.create_task(self._refresh_schema_index(account_id))
        except RuntimeError as exc:
            log.debug("Schema refresh scheduling failed for %s: %s", account_id, exc)

    async def _refresh_schema_index(self, account_id: str) -> None:
        try:
            async with AsyncSessionLocal() as session:
                gateway = await SalesforceGateway.for_tenant(str(account_id), session)
                indexer = SalesforceSchemaIndexer(session)
                await indexer.ensure_quick_index(
                    org_id=str(account_id), gateway=gateway
                )
        except Exception as exc:  # noqa: BLE001
            log.debug("Schema refresh failed for %s: %s", account_id, exc)

    def _research_cache_key(self, account_id: str, prompt: str) -> str:
        return make_cache_key(
            "ai_reports_research",
            {"org_id": account_id, "prompt": prompt},
        )

    async def _get_cached_research(
        self, account_id: str, prompt: str
    ) -> Optional[ResearchArtifacts]:
        if not prompt or AI_REPORTS_RESEARCH_CACHE_TTL <= 0:
            return None
        key = self._research_cache_key(account_id, prompt)
        cached = await self._research_cache.get(key)
        if not cached:
            return None
        try:
            return ResearchArtifacts.model_validate(cached)
        except ValidationError:
            return None

    async def _cache_research(
        self, account_id: str, prompt: str, research: ResearchArtifacts
    ) -> None:
        if not prompt or AI_REPORTS_RESEARCH_CACHE_TTL <= 0:
            return
        key = self._research_cache_key(account_id, prompt)
        await self._research_cache.set(key, research.model_dump())

    @staticmethod
    def _extract_object_refs(expr: str, primary: Optional[str]) -> Set[str]:
        if not expr:
            return set()
        refs = set(re.findall(r"([A-Za-z0-9_]+)\.", expr))
        if not refs and primary:
            refs.add(primary)
        return {ref.lower() for ref in refs}

    def _plan_uses_secondary_objects(self, plan: SalesforcePlan) -> bool:
        if not plan.objects or len(plan.objects) <= 1:
            return False
        primary = plan.objects[0]
        secondary = {obj.lower() for obj in plan.objects[1:]}

        for field in plan.fields or []:
            if self._extract_object_refs(field, primary) & secondary:
                return True
        for field in plan.group_by or []:
            if self._extract_object_refs(field, primary) & secondary:
                return True
        for order in plan.order_by or []:
            if self._extract_object_refs(order.field, primary) & secondary:
                return True
        for filt in plan.filters or []:
            if self._extract_object_refs(filt.field, primary) & secondary:
                return True
        return False

    def _prune_unused_joins(self, plan: SalesforcePlan) -> SalesforcePlan:
        if not plan.objects:
            return plan
        primary = plan.objects[0]
        if len(plan.objects) <= 1:
            plan.objects = [primary]
            plan.joins = []
            return plan

        used_objects = {primary.lower()}

        def add_refs(expr: str) -> None:
            used_objects.update(self._extract_object_refs(expr, primary))

        for field in plan.fields or []:
            add_refs(field)
        for group in plan.group_by or []:
            add_refs(group)
        for order in plan.order_by or []:
            add_refs(order.field)
        for filt in plan.filters or []:
            add_refs(filt.field)

        kept_objects = [
            obj for obj in plan.objects if obj == primary or obj.lower() in used_objects
        ]
        if not kept_objects:
            kept_objects = [primary]
        plan.objects = kept_objects

        if plan.joins:
            kept = {obj.lower() for obj in kept_objects}
            plan.joins = [
                join
                for join in plan.joins
                if join.from_object.lower() in kept and join.to_object.lower() in kept
            ]
        return plan

    def _plan_requires_critic(self, plan: SalesforcePlan) -> bool:
        # 1. Mandatory Critic if the plan is hollow
        if not plan.objects:
            return True
        if not plan.fields:
            # Empty fields array is suspicious - critic should validate
            log.info("Critic REQUIRED: Plan is empty or missing fields.")
            return True

        # 2. Trust the Planner if it provided Joins + Aggregates
        # (This indicates high-level reasoning is already complete)
        has_aggregates = any("(" in f for f in plan.fields)
        if len(plan.objects) > 1 and plan.joins and has_aggregates:
            log.info(
                "Critic SKIPPED: Planner provided valid structural join and aggregates."
            )
            return False

        # 3. Standard complexity checks
        if len(plan.joins or []) > 1:
            return True
        if len(plan.group_by or []) > 1:
            return True
        if self._plan_uses_secondary_objects(plan):
            return True
        # Check if multiple objects are referenced in fields but no joins defined
        if plan.fields and not plan.joins:
            primary = plan.objects[0] if plan.objects else None
            referenced_objects = set()
            for field in plan.fields:
                if "." in field:
                    obj = field.split(".")[0]
                    # Skip aggregate functions like SUM(Opportunity.Amount)
                    if "(" in obj:
                        continue
                    referenced_objects.add(obj)
            # If fields reference objects other than the primary, joins are needed
            if primary and len(referenced_objects - {primary}) > 0:
                return True
        return False

    def _is_simple_request(
        self, prompt: str, research: Optional[ResearchArtifacts]
    ) -> bool:
        if research and research.related_objects:
            return len(research.related_objects) <= 1
        return len(_tokenize(prompt)) <= 12

    def _planner_thinking_level(
        self, prompt: str, research: Optional[ResearchArtifacts]
    ) -> Optional[str]:
        base = AI_REPORTS_THINKING_LEVELS.get("planner", "MEDIUM")
        if not isinstance(base, str):
            return base
        normalized = base.strip().upper() or "MEDIUM"
        if normalized in {"MINIMAL", "LOW"}:
            return normalized
        if self._is_simple_request(prompt, research):
            return "LOW"
        return normalized

    async def _get_prompt_embedding(self, text: str) -> Optional[List[float]]:
        normalized = (text or "").strip()
        if not normalized:
            return None
        if self._embeddings_disabled:
            return None
        try:
            model = self._ensure_embedding_model()
            if model is None:
                return None
        except Exception as exc:  # noqa: BLE001
            if self._embeddings_disabled:
                return None
            log.debug("Embedding model unavailable: %s", exc)
            return None
        try:
            response = await asyncio.to_thread(model.get_embeddings, [normalized])
        except Exception as exc:  # noqa: BLE001
            if self._maybe_disable_embeddings(exc):
                return None
            log.warning("Vertex embedding request failed: %s", exc)
            return None
        if not response:
            return None
        values = getattr(response[0], "values", None)
        if not values:
            return None
        return list(values)

    async def _build_semantic_context(
        self,
        *,
        org_id: str,
        repo: SalesforceSchemaIndexRepo,
        prompt: str,
    ) -> Optional[SchemaPromptContext]:
        vector = await self._get_prompt_embedding(prompt)
        if not vector:
            return None
        object_limit = max(1, min(3, AI_REPORTS_MAX_OBJECTS))
        field_limit = max(5, min(15, AI_REPORTS_MAX_FIELDS))
        top_objects = await repo.find_similar_objects(org_id, vector, limit=object_limit)
        if not top_objects:
            return None
        object_filter = [obj.object_api for obj in top_objects if obj.object_api]
        full_fields_by_object: Dict[str, List[Any]] = {}
        for obj in top_objects:
            api_name = getattr(obj, "object_api", None)
            if not api_name:
                continue
            try:
                full_fields_by_object[api_name] = await repo.list_fields(org_id, api_name)
            except Exception as exc:  # noqa: BLE001
                log.debug("Failed to fetch full fields for %s: %s", api_name, exc)
        top_fields = await repo.find_similar_fields(
            org_id,
            vector,
            limit=field_limit,
            object_filter=object_filter or None,
        )
        return self._build_context_from_semantic_results(
            top_objects=top_objects,
            top_fields=top_fields,
            prompt=prompt,
            full_fields_by_object=full_fields_by_object,
        )

    def _build_context_from_semantic_results(
        self,
        *,
        top_objects: List[Any],
        top_fields: List[Any],
        prompt: str,
        full_fields_by_object: Optional[Dict[str, List[Any]]] = None,
    ) -> Optional[SchemaPromptContext]:
        if not top_objects:
            return None
        tokens = _tokenize(prompt)
        objects_by_name: Dict[str, SchemaObject] = {}
        prompt_objects: List[PromptObject] = []
        fields_by_object: Dict[str, List[Any]] = defaultdict(list)
        full_fields_map = full_fields_by_object or {}
        for field_row in top_fields or []:
            if field_row and getattr(field_row, "object_api", None):
                fields_by_object[field_row.object_api].append(field_row)
        for obj in top_objects:
            api_name = getattr(obj, "object_api", None)
            if not api_name:
                continue
            top_field_names = {
                getattr(field_row, "field_api", None)
                for field_row in fields_by_object.get(api_name, [])
                if getattr(field_row, "field_api", None)
            }
            fields_source = full_fields_map.get(api_name) or fields_by_object.get(
                api_name, []
            )
            schema_obj = SchemaObject(
                name=api_name,
                label=getattr(obj, "label", None) or api_name,
                label_plural=getattr(obj, "label_plural", None)
                or getattr(obj, "label", None)
                or api_name,
            )
            sample_fields: List[str] = []
            for field_row in fields_source:
                field_api = getattr(field_row, "field_api", None)
                if not field_api:
                    continue
                field_accessible = bool(getattr(field_row, "accessible", True))
                if not field_accessible:
                    continue
                try:
                    picklist_values = json.loads(field_row.picklist_json or "[]")
                except Exception:  # noqa: BLE001
                    picklist_values = []
                schema_field = SchemaField(
                    name=field_api,
                    label=getattr(field_row, "label", None) or field_api,
                    data_type=getattr(field_row, "data_type", "") or "",
                    relationship_name=getattr(field_row, "relationship_api", None),
                    reference_to=[
                        ref.strip()
                        for ref in (getattr(field_row, "reference_to", "") or "").split(",")
                        if ref.strip()
                    ],
                    help_text=getattr(field_row, "help_text", None),
                    queryable=bool(getattr(field_row, "queryable", True)),
                    accessible=field_accessible,
                    picklist_values=picklist_values,
                    length=getattr(field_row, "length", None),
                )
                schema_obj.fields[schema_field.name.lower()] = schema_field
                if schema_field.relationship_name and schema_field.reference_to:
                    rel_key = schema_field.relationship_name.lower()
                    schema_obj.relationships.setdefault(rel_key, [])
                    schema_obj.relationships[rel_key].extend(schema_field.reference_to)
                if field_api in top_field_names and len(sample_fields) < 8:
                    sample_fields.append(field_api)
            if not sample_fields:
                for field_api in list(schema_obj.fields.keys())[:8]:
                    sample_fields.append(schema_obj.fields[field_api].name)
            objects_by_name[api_name.lower()] = schema_obj
            prompt_objects.append(
                PromptObject(
                    name=api_name,
                    label=schema_obj.label,
                    sample_fields=sample_fields,
                    relationship_hints=list(
                        {
                            rel
                            for rels in schema_obj.relationships.values()
                            for rel in rels
                        }
                    ),
                )
            )
        if not prompt_objects:
            return None
        return SchemaPromptContext(
            prompt_objects=prompt_objects,
            tokens=tokens,
            objects_by_name=objects_by_name,
        )

    async def build_schema_context(
        self,
        *,
        account_id: str,
        gateway: SalesforceGateway,
        prompt: str,
        db_session: Optional[AsyncSession] = None,
    ) -> SchemaPromptContext:
        if db_session:
            repo = SalesforceSchemaIndexRepo(db_session)
            try:
                status = {}
                index_status = None
                has_index = False
                try:
                    status = await repo.status(
                        account_id, ttl_hours=AI_SCHEMA_INDEX_TTL_HOURS
                    )
                    index_status = status.get("index_status")
                    has_index = (status.get("objects") or 0) > 0 and (
                        status.get("fields") or 0
                    ) > 0
                except Exception as exc:  # noqa: BLE001
                    log.debug("Schema index status unavailable: %s", exc)

                if has_index:
                    if index_status == "stale" and not status.get("in_progress"):
                        self._schedule_schema_refresh(account_id)
                    if index_status in {"quick", "full"}:
                        semantic_context = await self._build_semantic_context(
                            org_id=account_id,
                            repo=repo,
                            prompt=prompt,
                        )
                        if semantic_context:
                            return semantic_context
                    indexed_context = await self._build_context_from_index(
                        org_id=account_id,
                        repo=repo,
                        prompt=prompt,
                    )
                    if indexed_context:
                        return indexed_context
            except Exception as exc:  # noqa: BLE001
                log.warning(
                    "Schema index fallback for %s due to error: %s", account_id, exc
                )
            try:
                indexer = SalesforceSchemaIndexer(db_session)
                await indexer.ensure_quick_index(org_id=account_id, gateway=gateway)
                semantic_context = await self._build_semantic_context(
                    org_id=account_id,
                    repo=repo,
                    prompt=prompt,
                )
                if semantic_context:
                    return semantic_context
                indexed_context = await self._build_context_from_index(
                    org_id=account_id,
                    repo=repo,
                    prompt=prompt,
                )
                if indexed_context:
                    return indexed_context
            except Exception as exc:  # noqa: BLE001
                log.warning(
                    "Schema index rebuild for %s failed: %s", account_id, exc
                )
        return await self.schema_index.build_prompt_context(
            account_id=account_id,
            gateway=gateway,
            prompt=prompt,
            max_objects=AI_REPORTS_MAX_OBJECTS,
            max_fields=AI_REPORTS_MAX_FIELDS,
        )

    async def _call_llm(
        self,
        client: VertexAIProvider,
        messages: List[Dict[str, Any]],
        schema_context: Optional[SchemaPromptContext] = None,
        thinking_level: Optional[str] = None,
        include_thoughts: Optional[bool] = None,
        schema: Optional[Dict[str, Any]] = None,
        max_output_tokens: Optional[int] = None,
        retry_on_parse_error: bool = True,
        timeout_s: Optional[int] = None,
    ) -> Dict[str, Any]:
        """Call Vertex AI with the given messages.

        Args:
            client: The VertexAIProvider client
            messages: List of message dicts with 'role' and 'content'
            schema_context: Optional schema context for caching
            thinking_level: Optional Gemini 3 thinking level (MINIMAL/LOW/MEDIUM/HIGH)
            include_thoughts: Optional flag to request Gemini thought trace
            schema: Optional JSON schema for structured output
            max_output_tokens: Optional override for generation max tokens
            retry_on_parse_error: Whether to retry once on parse errors
            timeout_s: Optional per-call timeout override

        Returns:
            Parsed JSON response dict

        Raises:
            PlannerUnavailable: If the LLM call fails
        """
        system_parts: List[str] = []
        user_parts: List[str] = []
        for msg in messages:
            role = (msg.get("role") or "user").lower()
            content = msg.get("content")
            if not content:
                continue
            if role == "system":
                system_parts.append(str(content))
            else:
                user_parts.append(str(content))
        combined_prompt = "\n\n".join(system_parts + user_parts).strip()

        # Get or create cache if schema_context is provided
        cache_name = ""
        if schema_context:
            cache_name = self._get_or_create_cache(schema_context, client)

        def _parse_llm_response(resp: Dict[str, Any]) -> Dict[str, Any]:
            """Strictly parse LLM response to JSON dict."""
            def _unwrap_value_list(payload: Any) -> Optional[Dict[str, Any]]:
                if not isinstance(payload, list) or len(payload) != 1:
                    return None
                if isinstance(payload[0], dict):
                    return payload[0]
                return None

            def _normalize_payload(payload: Any) -> Dict[str, Any]:
                if isinstance(payload, list):
                    unwrapped = _unwrap_value_list(payload)
                    if unwrapped:
                        return unwrapped
                    return {"value": payload}
                if isinstance(payload, dict):
                    if set(payload.keys()) == {"value"} and isinstance(
                        payload.get("value"), list
                    ):
                        unwrapped = _unwrap_value_list(payload["value"])
                        if unwrapped:
                            return unwrapped
                    return payload
                return {"value": payload}

            def _parse_text_payload(raw_text: str) -> Dict[str, Any]:
                if not isinstance(raw_text, str) or not raw_text.strip():
                    return {"_parse_error": True, "_raw": ""}

                cleaned = raw_text.strip()
                if cleaned.startswith("```"):
                    lines = cleaned.split("\n", 1)
                    if len(lines) > 1:
                        cleaned = lines[1]
                    else:
                        cleaned = ""
                if cleaned.endswith("```"):
                    cleaned = cleaned[:-3]
                cleaned = cleaned.strip()

                if not cleaned:
                    return {"_parse_error": True, "_raw": raw_text[:500]}

                try:
                    parsed = json.loads(cleaned)
                    return _normalize_payload(parsed)
                except json.JSONDecodeError:
                    pass

                repaired = _repair_json(cleaned)
                if repaired:
                    try:
                        parsed = json.loads(repaired)
                        log.info("JSON repair successful for LLM response")
                        return _normalize_payload(parsed)
                    except json.JSONDecodeError:
                        pass

                def _truncate_on_repeat(text: str) -> Optional[str]:
                    markers = (
                        "plan",
                        "clarifying_questions",
                        "risks",
                    )
                    for key in markers:
                        pattern = re.compile(r'\{\s*"' + re.escape(key) + r'"\s*:', re.I)
                        matches = list(pattern.finditer(text))
                        if len(matches) > 1:
                            return text[: matches[1].start()].strip()
                    return None

                truncated = _truncate_on_repeat(cleaned)
                if truncated:
                    repaired = _repair_json(truncated)
                    if repaired:
                        try:
                            parsed = json.loads(repaired)
                            log.info("Salvaged repeated JSON payload")
                            return _normalize_payload(parsed)
                        except json.JSONDecodeError:
                            pass

                def _extract_balanced_objects(text: str) -> List[str]:
                    objs = []
                    depth = 0
                    start = None
                    in_str = False
                    escape = False
                    for idx, ch in enumerate(text):
                        if in_str:
                            if escape:
                                escape = False
                            elif ch == "\\":
                                escape = True
                            elif ch == '"':
                                in_str = False
                            continue
                        if ch == '"':
                            in_str = True
                            continue
                        if ch == "{":
                            if depth == 0:
                                start = idx
                            depth += 1
                            continue
                        if ch == "}":
                            if depth > 0:
                                depth -= 1
                                if depth == 0 and start is not None:
                                    objs.append(text[start : idx + 1])
                                    start = None
                    return objs

                candidates: List[str] = []
                start_obj = cleaned.find("{")
                start_arr = cleaned.find("[")
                if start_arr != -1 and (start_obj == -1 or start_arr < start_obj):
                    end = cleaned.rfind("]")
                    if end > start_arr:
                        candidates.append(cleaned[start_arr : end + 1])
                elif start_obj != -1:
                    end = cleaned.rfind("}")
                    if end > start_obj:
                        candidates.append(cleaned[start_obj : end + 1])

                candidates.extend(_extract_balanced_objects(cleaned))
                for payload in candidates:
                    try:
                        parsed = json.loads(payload)
                    except json.JSONDecodeError:
                        repaired = _repair_json(payload)
                        if not repaired:
                            continue
                        try:
                            parsed = json.loads(repaired)
                            log.info("JSON repair successful for LLM response")
                        except json.JSONDecodeError:
                            continue
                    return _normalize_payload(parsed)

                return {"_parse_error": True, "_raw": raw_text[:500]}

            value = resp.get("value")
            if isinstance(value, dict):
                if "raw_text" in value and len(value.keys()) == 1:
                    return _parse_text_payload(value.get("raw_text") or "")
                return _normalize_payload(value)
            if isinstance(value, list):
                return _normalize_payload(value)

            raw_text = resp.get("raw_text") or resp.get("value") or ""
            return _parse_text_payload(raw_text)

        # First attempt
        if timeout_s is None:
            timeout_s = AI_REPORTS_TIMEOUT_S
        try:
            response = await asyncio.to_thread(
                client.generate_json,
                combined_prompt,
                schema=schema,
                cache_name=cache_name,
                thinking_level=thinking_level,
                include_thoughts=include_thoughts,
                max_output_tokens=max_output_tokens,
                timeout_s=timeout_s,
            )
        except Exception as exc:
            log.error("Vertex AI call failed: %s", exc)
            raise PlannerUnavailable(f"Vertex AI error: {exc}") from exc

        result = _parse_llm_response(response)

        # If parse failed, retry once with strict instructions
        if result.get("_parse_error") and retry_on_parse_error:
            log.warning(
                "LLM returned non-JSON; retrying with strict prompt. Raw snippet: %s",
                result.get("_raw", "")[:200],
            )
            retry_prompt = (
                f"{combined_prompt}\n\n"
                "IMPORTANT: Return ONLY valid JSON. No prose, no markdown, no code fences. "
                "Match the required schema exactly."
            )
            try:
                retry_response = await asyncio.to_thread(
                    client.generate_json,
                    retry_prompt,
                    schema=schema,
                    cache_name=cache_name,
                    thinking_level=thinking_level,
                    include_thoughts=include_thoughts,
                    max_output_tokens=max_output_tokens,
                    timeout_s=timeout_s,
                )
                result = _parse_llm_response(retry_response)
                if result.get("_parse_error"):
                    log.error(
                        "LLM retry also failed to produce valid JSON: %s",
                        result.get("_raw", "")[:200],
                    )
            except Exception as retry_exc:
                log.error("LLM retry call failed: %s", retry_exc)
                # Keep original parse_error result

        return result

    async def _run_researcher(
        self,
        *,
        prompt: str,
        schema_context: SchemaPromptContext,
    ) -> ResearchArtifacts:
        model_name = AI_REPORTS_RESEARCHER_MODEL or AI_REPORTS_PLANNER_MODEL
        client = self._get_client(model_name)
        messages = [
            {"role": "system", "content": self._researcher_system_prompt()},
            {
                "role": "user",
                "content": self._build_researcher_prompt(
                    prompt=prompt,
                    schema_context=schema_context,
                ),
            },
        ]
        # Researcher is intent-only; disable thinking for speed and consistency.
        thinking_level = None
        max_output_tokens = AI_REPORTS_RESEARCHER_MAX_OUTPUT_TOKENS
        if max_output_tokens is None or max_output_tokens <= 0:
            max_output_tokens = 1024
        try:
            raw = await self._call_llm(
                client,
                messages,
                thinking_level=thinking_level,
                schema=RESEARCHER_SCHEMA,
                max_output_tokens=max_output_tokens,
                retry_on_parse_error=False,
                timeout_s=AI_REPORTS_RESEARCHER_TIMEOUT_S,
            )
        except PlannerUnavailable as exc:
            log.warning(
                "Researcher model unavailable (%s); using heuristic discovery.", exc
            )
            return self._fallback_research(prompt=prompt, schema_context=schema_context)
        if raw.get("_parse_error"):
            log.warning("Researcher parse error: %s", raw.get("_raw"))
            return self._fallback_research(prompt=prompt, schema_context=schema_context)
        try:
            artifacts = ResearchArtifacts.model_validate(raw)
        except ValidationError as exc:
            log.warning("Researcher produced invalid payload: %s", exc)
            log.debug(
                "Researcher raw response: %s", json.dumps(raw, indent=2, default=str)
            )
            return self._fallback_research(prompt=prompt, schema_context=schema_context)

        artifacts.clarifying_questions = artifacts.clarifying_questions[:3]
        artifacts.related_objects = artifacts.related_objects[:5]
        artifacts.suggested_fields = artifacts.suggested_fields[:8]
        artifacts.suggested_filters = artifacts.suggested_filters[:5]
        artifacts.common_filters = artifacts.common_filters[:6] or [
            "CloseDate = THIS_QUARTER",
            "IsClosed = false",
        ]
        artifacts.example_plans = artifacts.example_plans[:3]
        artifacts.macros = self.list_macros()
        return artifacts

    async def _run_researcher_consolidated(
        self,
        *,
        account_id: str,
        prompt: str,
        schema_context: SchemaPromptContext,
    ) -> ResearchArtifacts:
        key = f"{account_id}:{prompt}"
        async with self._research_lock:
            task = self._in_flight_research.get(key)
            if task is None or task.done():
                task = asyncio.create_task(
                    self._run_researcher(prompt=prompt, schema_context=schema_context)
                )
                self._in_flight_research[key] = task
                log.debug("Started consolidated research task for %s", key)
            else:
                log.info("Piggybacking on in-flight research for %s", key)
        try:
            return await asyncio.shield(task)
        finally:
            async with self._research_lock:
                current = self._in_flight_research.get(key)
                if current is task and task.done():
                    self._in_flight_research.pop(key, None)

    async def _run_planner(
        self,
        *,
        prompt: str,
        schema_context: SchemaPromptContext,
        research: Optional[ResearchArtifacts],
        tier: Optional[str],
        prior_warnings: Optional[List[str]] = None,
        conversation: Optional[List[Dict[str, str]]] = None,
        clarifier_answers: Optional[Dict[str, str]] = None,
    ) -> PlannerDraft:
        model_name = self.select_model_for_tier(tier)
        client = self._get_client(model_name)
        thinking_level = self._planner_thinking_level(prompt, research)
        max_output_tokens = (
            AI_REPORTS_PLANNER_MAX_OUTPUT_TOKENS
            if AI_REPORTS_PLANNER_MAX_OUTPUT_TOKENS > 0
            else None
        )
        messages = [
            {"role": "system", "content": self._planner_system_prompt()},
            {
                "role": "user",
                "content": self._build_planner_prompt(
                    prompt=prompt,
                    schema_context=schema_context,
                    research=research,
                    row_cap=self.tier_row_cap(tier),
                    preview_limit=AI_REPORTS_PREVIEW_LIMIT,
                    prior_warnings=prior_warnings,
                    conversation=conversation,
                    clarifier_answers=clarifier_answers,
                ),
            },
        ]
        try:
            raw = await self._call_llm(
                client,
                messages,
                schema_context=schema_context,
                thinking_level=thinking_level,
                schema=PLANNER_SCHEMA,
                max_output_tokens=max_output_tokens,
                timeout_s=AI_REPORTS_PLANNER_TIMEOUT_S,
            )
        except PlannerUnavailable as primary_exc:
            fallback_model = AI_REPORTS_RESEARCHER_MODEL or None
            if fallback_model and fallback_model != model_name:
                log.info(
                    "Planner timeout on %s; retrying with fallback model %s",
                    model_name,
                    fallback_model,
                )
                fallback_client = self._get_client(fallback_model)
                try:
                    raw = await self._call_llm(
                        fallback_client,
                        messages,
                        schema_context=schema_context,
                        thinking_level=thinking_level,
                        schema=PLANNER_SCHEMA,
                        max_output_tokens=max_output_tokens,
                        timeout_s=AI_REPORTS_PLANNER_TIMEOUT_S,
                    )
                except PlannerUnavailable as secondary_exc:
                    log.error(
                        "Fallback planner model %s also unavailable: %s",
                        fallback_model,
                        secondary_exc,
                    )
                    raise primary_exc
            else:
                raise primary_exc
        if raw.get("_parse_error"):
            log.error(
                "Planner parse error for prompt '%s': %s", prompt, raw.get("_raw")
            )
            raise PlannerUnavailable("Planner returned unparseable response")
        # Debug: log raw planner output before validation
        log.debug(
            "Planner raw output: %s",
            json.dumps(raw, indent=2, default=str)[:2000],
        )
        try:
            draft = PlannerDraft.model_validate(raw)
        except ValidationError as exc:
            log.error("Planner output invalid: %s", exc)
            log.debug(
                "Planner raw response: %s", json.dumps(raw, indent=2, default=str)
            )
            raise PlannerUnavailable("Planner produced invalid response")
        # Debug: log parsed plan structure
        log.info(
            "Planner draft: objects=%r, fields=%d, filters=%d",
            draft.plan.objects,
            len(draft.plan.fields),
            len(draft.plan.filters),
        )

        draft.trace = draft.trace[:8]
        # Generate fallback trace if the model didn't provide one
        if not draft.trace:
            fallback_trace = []
            if draft.plan.objects:
                fallback_trace.append(f"Primary object -> {draft.plan.objects[0]}")
            if draft.plan.fields:
                fallback_trace.append(
                    f"Selected fields -> {', '.join(draft.plan.fields[:4])}"
                    + ("..." if len(draft.plan.fields) > 4 else "")
                )
            if draft.plan.filters:
                filter_summary = "; ".join(
                    f"{f.field} {f.op} {f.value}"
                    for f in draft.plan.filters[:2]
                )
                fallback_trace.append(f"Filters -> {filter_summary}")
            if draft.plan.group_by:
                fallback_trace.append(f"Grouped by -> {', '.join(draft.plan.group_by)}")
            draft.trace = fallback_trace[:8]
        normalized_clarifiers: List[PlanClarifier] = []
        for idx, clarifier in enumerate(draft.clarifiers[:4]):
            if not clarifier.id:
                clarifier.id = f"clarifier_{idx+1}"
            clarifier.options = clarifier.options[:3]
            for option in clarifier.options:
                if option.patch:
                    option.patch = _ensure_patch_summary(option.patch)
            normalized_clarifiers.append(clarifier)
        draft.clarifiers = normalized_clarifiers
        return draft

    async def _run_critic(
        self,
        *,
        prompt: str,
        schema_context: SchemaPromptContext,
        research: Optional[ResearchArtifacts],
        plan: SalesforcePlan,
        tier: Optional[str],
    ) -> CriticFeedback:
        model_name = (
            AI_REPORTS_CRITIC_MODEL
            or AI_REPORTS_RESEARCHER_MODEL
            or self.select_model_for_tier(tier)
        )
        client = self._get_client(model_name)
        messages = [
            {"role": "system", "content": self._critic_system_prompt()},
            {
                "role": "user",
                "content": self._build_critic_prompt(
                    prompt=prompt,
                    schema_context=schema_context,
                    research=research,
                    plan=plan,
                    row_cap=self.tier_row_cap(tier),
                ),
            },
        ]
        try:
            raw = await self._call_llm(
                client,
                messages,
                schema_context=schema_context,
                thinking_level=AI_REPORTS_THINKING_LEVELS.get("critic", "LOW"),
                include_thoughts=True,
                schema=CRITIC_SCHEMA,
            )
        except PlannerUnavailable:
            fallback_model = AI_REPORTS_RESEARCHER_MODEL
            if fallback_model and fallback_model != model_name:
                log.info(
                    "Critic timeout on %s; retrying with fallback model %s",
                    model_name,
                    fallback_model,
                )
                fallback_client = self._get_client(fallback_model)
                try:
                    raw = await self._call_llm(
                        fallback_client,
                        messages,
                        schema_context=schema_context,
                        thinking_level=AI_REPORTS_THINKING_LEVELS.get("critic", "LOW"),
                        include_thoughts=True,
                        schema=CRITIC_SCHEMA,
                    )
                except PlannerUnavailable as secondary_exc:
                    log.warning(
                        "Critic fallback model %s also unavailable: %s",
                        fallback_model,
                        secondary_exc,
                    )
                    return self._fallback_critic(
                        plan=plan, schema_context=schema_context
                    )
            else:
                return self._fallback_critic(plan=plan, schema_context=schema_context)
        if raw.get("_parse_error"):
            log.warning("Critic parse error: %s", raw.get("_raw"))
            return self._fallback_critic(plan=plan, schema_context=schema_context)
        try:
            feedback = CriticFeedback.model_validate(raw)
        except ValidationError as exc:
            log.warning("Critic output invalid: %s", exc)
            log.debug("Critic raw response: %s", json.dumps(raw, indent=2, default=str))
            return self._fallback_critic(plan=plan, schema_context=schema_context)
        feedback.risks = feedback.risks[:5] or self._fallback_risks(
            plan, schema_context
        )
        normalized_alternatives: List[PlanAlternative] = []
        for alt in feedback.alternatives[:4]:
            alt.patch = _ensure_patch_summary(alt.patch)
            normalized_alternatives.append(alt)
        if not normalized_alternatives:
            normalized_alternatives = self._fallback_alternatives(plan)
        feedback.alternatives = normalized_alternatives
        return feedback

    def _fallback_research(
        self,
        *,
        prompt: str,
        schema_context: SchemaPromptContext,
    ) -> ResearchArtifacts:
        tokens = {tok for tok in schema_context.tokens}
        clarifiers: List[str] = []
        if "quarter" not in tokens:
            clarifiers.append(
                "Which date range should we use? (THIS_QUARTER, LAST_90_DAYS)"
            )
        if "pipeline" in tokens or "open" in tokens:
            clarifiers.append("Include only open pipeline stages? (Yes/No)")
        if not clarifiers:
            clarifiers.append("Any specific stage or owner to focus on?")

        related_objects = [
            SuggestedObject(name=obj.name, why=f"Matches {obj.label}")
            for obj in schema_context.prompt_objects[:5]
        ]
        suggested_fields: List[SuggestedField] = []
        for obj in schema_context.prompt_objects[:3]:
            for field_name in obj.sample_fields[:3]:
                suggested_fields.append(
                    SuggestedField(
                        object=obj.name,
                        field=f"{obj.name}.{field_name}",
                        why=f"Common field on {obj.name}",
                    )
                )
        date_literal = self._default_close_date_literal(list(tokens))
        suggested_filters = [
            SuggestedFilter(
                expression=f"CloseDate = {date_literal}", why="Quarter-based snapshot"
            ),
            SuggestedFilter(expression="IsClosed = false", why="Open pipeline focus"),
        ]
        example_plans = [
            ExamplePlan(
                title=f"{schema_context.prompt_objects[0].name if schema_context.prompt_objects else 'Opportunity'} pipeline overview",
                summary="Show open pipeline amounts grouped by stage and owner.",
            )
        ]
        return ResearchArtifacts(
            clarifying_questions=clarifiers[:3],
            related_objects=related_objects,
            suggested_fields=suggested_fields[:8],
            suggested_filters=suggested_filters,
            common_filters=[f"CloseDate = {date_literal}", "IsClosed = false"],
            example_plans=example_plans,
            macros=self.list_macros(),
        )

    def _fallback_risks(
        self, plan: SalesforcePlan, schema_context: SchemaPromptContext
    ) -> List[str]:
        risks: List[str] = []
        if plan.limit is None:
            risks.append("Plan has no LIMIT; results may exceed org row caps.")
        if not plan.group_by and any("Stage" in field for field in plan.fields or []):
            risks.append("Consider grouping by StageName for pipeline summaries.")
        if not risks:
            risks.append("Review field-level security before running this plan.")
        return risks

    def _fallback_alternatives(self, plan: SalesforcePlan) -> List[PlanAlternative]:
        alternatives: List[PlanAlternative] = []
        if "Opportunity.StageName" not in (plan.fields or []):
            alternatives.append(
                PlanAlternative(
                    title="Add StageName field",
                    description="Expose opportunity stage to support grouping or filtering.",
                    patch=_ensure_patch_summary(
                        PlanPatch(
                            ops=[
                                PlanPatchOp(
                                    op="add",
                                    path="/fields/-",
                                    value="Opportunity.StageName",
                                )
                            ]
                        )
                    ),
                )
            )
        if plan.limit is None:
            alternatives.append(
                PlanAlternative(
                    title="Apply LIMIT 200",
                    description="Stay under preview caps before running the full query.",
                    patch=_ensure_patch_summary(
                        PlanPatch(
                            ops=[PlanPatchOp(op="replace", path="/limit", value=200)]
                        )
                    ),
                )
            )
        if not alternatives:
            alternatives.append(
                PlanAlternative(
                    title="Add Amount field",
                    description="Include Opportunity.Amount to show pipeline totals.",
                    patch=_ensure_patch_summary(
                        PlanPatch(
                            ops=[
                                PlanPatchOp(
                                    op="add",
                                    path="/fields/-",
                                    value="Opportunity.Amount",
                                )
                            ]
                        )
                    ),
                )
            )
        return alternatives

    def _fallback_critic(
        self,
        *,
        plan: SalesforcePlan,
        schema_context: SchemaPromptContext,
    ) -> CriticFeedback:
        return CriticFeedback(
            risks=self._fallback_risks(plan, schema_context),
            alternatives=self._fallback_alternatives(plan),
        )

    def _build_default_plan_payload(
        self,
        *,
        prompt: str,
        schema_context: SchemaPromptContext,
        research: Optional[ResearchArtifacts],
        ranked_fields: Optional[List[Dict[str, Any]]] = None,
        prompt_terms: Optional[List[str]] = None,
        org_id: Optional[str] = None,
        request_id: Optional[str] = None,
    ) -> Tuple[SalesforcePlan, List[str], List[str], List[str], List[PlanAlternative]]:
        primary = None
        if research and research.related_objects:
            primary = research.related_objects[0].name
        elif schema_context.prompt_objects:
            primary = schema_context.prompt_objects[0].name
        if not primary:
            primary = "Opportunity"

        plan = SalesforcePlan(
            objects=[primary],
            fields=[],
            joins=[],
            filters=[],
            order_by=[],
            group_by=[],
            limit=200,
        )

        def normalize_field(obj: str, field_name: str) -> str:
            if "." in field_name:
                return field_name
            return f"{obj}.{field_name}"

        candidates: List[str] = []
        if ranked_fields:
            for item in ranked_fields[:5]:
                field_ref = item.get("field_ref")
                if field_ref:
                    candidates.append(field_ref)
        if research and research.suggested_fields:
            for suggested in research.suggested_fields[:5]:
                obj = suggested.obj or primary
                candidates.append(normalize_field(obj, suggested.field))
        if not candidates and schema_context.prompt_objects:
            for obj in schema_context.prompt_objects:
                if obj.name == primary:
                    for field_name in obj.sample_fields[:4]:
                        candidates.append(normalize_field(primary, field_name))
                    break
        if not candidates:
            if primary.lower() == "opportunity":
                candidates = [
                    "Opportunity.Name",
                    "Opportunity.StageName",
                    "Opportunity.Amount",
                ]
            else:
                candidates = [f"{primary}.Name"]
        plan.fields = list(dict.fromkeys(candidates))

        # ─── AUTO-GENERATE JOINS FOR FALLBACK PLANNER ───
        # Detect fields referencing non-primary objects and generate joins
        common_lookups = {
            ("opportunity", "account"): ("AccountId", "Account"),
            ("opportunity", "campaign"): ("CampaignId", "Campaign"),
            ("contact", "account"): ("AccountId", "Account"),
            ("case", "account"): ("AccountId", "Account"),
            ("case", "contact"): ("ContactId", "Contact"),
            ("opportunitylineitem", "product2"): ("Product2Id", "Product2"),
            ("opportunitylineitem", "opportunity"): ("OpportunityId", "Opportunity"),
        }

        referenced_objects: Set[str] = set()
        for field in plan.fields:
            if "." in field:
                obj_name = field.split(".")[0].lower()
                if obj_name != primary.lower():
                    referenced_objects.add(obj_name)

        generated_joins: List[PlanJoin] = []
        for ref_obj in referenced_objects:
            key = (primary.lower(), ref_obj)
            if key in common_lookups:
                via, relationship = common_lookups[key]
                generated_joins.append(
                    PlanJoin(
                        from_object=primary,
                        to_object=ref_obj.title(),
                        via=via,
                        relationship=relationship,
                    )
                )
                # Add to objects list
                if ref_obj.title() not in plan.objects:
                    plan.objects.append(ref_obj.title())

        if generated_joins:
            plan.joins = generated_joins

        if primary.lower() == "opportunity":
            plan.order_by = [PlanOrder(field="Opportunity.CloseDate", dir="DESC")]
        elif primary.lower() == "case":
            plan.order_by = [PlanOrder(field="Case.CreatedDate", dir="DESC")]

        trace = [
            "Planner fallback due to model timeout or invalid planner output",
            f"Primary object -> {primary}",
            f"Selected fields -> {', '.join(plan.fields[:4])}",
        ]
        plan_trace = trace[:8]
        warnings = [
            "Used heuristic planner because the AI planner model timed out or returned an invalid response. "
            "Ensure the Gemini/Ollama service is reachable or select a smaller model.",
        ]
        alternatives = self._fallback_alternatives(plan)
        risks = self._fallback_risks(plan, schema_context)
        prompt_tokens = prompt_terms or _tokenize(prompt)
        plan = self._ensure_defaults(
            plan=plan,
            ranked_fields=ranked_fields or [],
            prompt_terms=prompt_tokens,
            org_id=org_id,
            request_id=request_id,
        )
        return plan, plan_trace, warnings, risks, alternatives

    def _researcher_system_prompt(self) -> str:
        return (
            "You are the Researcher agent for FoundryMatch's Salesforce AI Report Builder. "
            "Given user intent and schema hints, emit JSON with keys: "
            "clarifying_questions (max 3 short questions), "
            'related_objects (array of {"name","why"}), '
            'suggested_fields (array of {"object","field","why"}), '
            'suggested_filters (array of {"expression","why"}), '
            "common_filters (array of strings), "
            'example_plans (array of {"title","summary"}). '
            "Rank suggestions by usefulness, avoid duplicates, and never output SOQL."
        )

    def _planner_system_prompt(self) -> str:
        return (
            "You are a Senior RevOps Consultant building Salesforce query plans. "
            "Respond with a single JSON object. "
            "CRITICAL: The 'plan' object MUST come FIRST in your JSON response, before 'trace' or 'clarifiers'. "
            "This ensures the query logic is captured even if the response is truncated. "
            "JSON keys in order: plan, trace, clarifiers. "
            "The plan schema:\n"
            "{\n"
            '  "objects": ["PrimaryObject", "..."],\n'
            '  "fields": ["Opportunity.Name", "Account.Name"],\n'
            '  "joins": [{"from":"Opportunity","to":"Account","via":"AccountId","relationship":"Account"}],\n'
            '  "filters": [{"field":"Opportunity.CloseDate","op":"DATE_LITERAL","value":"THIS_QUARTER"}],\n'
            '  "orderBy": [{"field":"Opportunity.CloseDate","dir":"DESC"}],\n'
            '  "groupBy": ["Opportunity.StageName"],\n'
            '  "limit": 200\n'
            "}\n"
            "Only use fields that appear in the provided schema hints or research summaries. "
            "STRICT AGGREGATION RULE: If groupBy has any fields, you must not include raw metric fields alongside aggregates. "
            "INCORRECT: ['StageName', 'Amount', 'SUM(Amount)']  CORRECT: ['StageName', 'SUM(Amount)']. "
            "Every non-grouped field must be an aggregate (SUM/AVG/COUNT/MAX/MIN); never select raw and aggregate together. "
            "Win rate must use IsWon (e.g., AVG(IsWon) AS win_rate). "
            "If the user asks for records without related children (e.g., accounts with no opportunities), use an anti-join/NOT EXISTS pattern instead of returning unrelated fields. "
            "REQUIRED: The trace array MUST contain 2-3 brief bullets explaining your logic (e.g., 'Primary object: Opportunity', 'Filter: open pipeline this quarter'). "
            "REQUIRED: The plan.objects array MUST contain at least one Salesforce object (e.g., ['Opportunity']). Never return an empty objects array. "
            "MANDATORY: The plan.fields array MUST contain every metric/field mentioned in the query. NEVER return an empty fields array. "
            "MANDATORY: If fields reference multiple objects (e.g., Account.Name with Opportunity.Amount), you MUST populate the joins array with the correct relationship. "
            "Clarifiers array is optional; if needed, use simple {\"id\",\"question\",\"default\"} format. "
            "Always include limit (default 200). Respect row/field caps. Return JSON only."
        )

    def _critic_system_prompt(self) -> str:
        return (
            "You are the Critic agent reviewing a Salesforce query plan. "
            "Return JSON with keys: risks (>=1 concise bullets) and alternatives (>=1). "
            'Each alternative must include {"title","description","patch"} '
            'where patch is JSON with "ops" array of JSON-Patch operations. '
            "Focus on field-level security, row limits, missing clarifier answers, and safer defaults. "
            "Explicitly flag SOQL traps: SELECT * or wildcard fields, mixing raw + aggregate "
            "fields when groupBy is used, OFFSET > 2000, semi-join/subquery limits in WHERE "
            "(IN/NOT IN subselects), missing filters or LIMIT on high-volume objects "
            "(Task, Event, EmailMessage, ActivityHistory), invalid field/relationship names, "
            "and cross-object fields without joins. "
            "Never change plan semantics without explaining why."
        )

    def _build_researcher_prompt(
        self,
        *,
        prompt: str,
        schema_context: SchemaPromptContext,
    ) -> str:
        ctx_block = schema_context.to_prompt_block()
        token_line = ", ".join(schema_context.tokens[:10])
        return (
            f"User intent:\n{prompt}\n\n"
            f"Schema hints:\n{ctx_block}\n\n"
            f"Keywords: {token_line}\n"
            "Respond with ranked JSON suggestions only."
        )

    def _build_planner_prompt(
        self,
        *,
        prompt: str,
        schema_context: SchemaPromptContext,
        research: Optional[ResearchArtifacts],
        row_cap: int,
        preview_limit: int,
        prior_warnings: Optional[List[str]] = None,
        conversation: Optional[List[Dict[str, str]]] = None,
        clarifier_answers: Optional[Dict[str, str]] = None,
    ) -> str:
        ctx_block = schema_context.to_prompt_block()
        warnings_block = ""
        if prior_warnings:
            warnings_block = (
                "Previous issues:\n- " + "\n- ".join(prior_warnings[:5]) + "\n\n"
            )
        convo_tail = ""
        if conversation:
            convo_lines = []
            for turn in conversation[-4:]:
                role = (turn.get("role") or "user").upper()
                content = turn.get("content", "")
                # Sanitize conversation content to prevent prompt injection
                content = content[:1000]  # Limit per-turn content
                content = content.replace("```", "")  # Prevent code block injection
                convo_lines.append(f"{role}: {content}")
            convo_tail = (
                "Recent conversation context:\n" + "\n".join(convo_lines) + "\n\n"
            )
        clarifier_block = ""
        if clarifier_answers:
            # Sanitize user input to prevent prompt injection attacks
            # Escape characters that could be used to break out of the context
            def sanitize_clarifier_value(val: str) -> str:
                if not isinstance(val, str):
                    val = str(val)
                # Limit length and remove control characters
                val = val[:500]  # Reasonable max length for a clarifier value
                # Replace potentially dangerous patterns
                val = val.replace("\n", " ").replace("\r", " ")
                val = val.replace("```", "")  # Prevent code block injection
                return val

            formatted = ", ".join(
                f"{key}={sanitize_clarifier_value(value)}"
                for key, value in clarifier_answers.items()
            )
            clarifier_block = f"Clarifier answers from user: {formatted}\n\n"
        research_json = (
            json.dumps(research.model_dump(exclude_none=True), ensure_ascii=False)
            if research
            else "{}"
        )
        return (
            f"{warnings_block}"
            f"{convo_tail}"
            f"{clarifier_block}"
            f"User intent:\n{prompt}\n\n"
            f"Schema hints:\n{ctx_block}\n\n"
            f"Research summary JSON:\n{research_json}\n\n"
            f"Constraints: preview_limit={preview_limit} rows, hard_row_cap={row_cap}. "
            "Always produce deterministic JSON exactly matching the required schema."
        )

    def _build_critic_prompt(
        self,
        *,
        prompt: str,
        schema_context: SchemaPromptContext,
        research: Optional[ResearchArtifacts],
        plan: SalesforcePlan,
        row_cap: int,
    ) -> str:
        ctx_block = schema_context.to_prompt_block()
        research_json = (
            json.dumps(research.model_dump(exclude_none=True), ensure_ascii=False)
            if research
            else "{}"
        )
        plan_json = json.dumps(plan.model_dump(by_alias=True), ensure_ascii=False)
        return (
            f"User intent:\n{prompt}\n\n"
            f"Schema hints:\n{ctx_block}\n\n"
            f"Research summary JSON:\n{research_json}\n\n"
            f"Proposed plan JSON:\n{plan_json}\n\n"
            f"Row cap: {row_cap}. Validate caps, FLS, and risky defaults."
        )

    async def plan_query(
        self,
        *,
        account_id: str,
        tier: Optional[str],
        prompt: str,
        gateway: SalesforceGateway,
        prior_warnings: Optional[List[str]] = None,
        conversation: Optional[List[Dict[str, str]]] = None,
        schema_context: Optional[SchemaPromptContext] = None,
        clarifier_answers: Optional[Dict[str, str]] = None,
        research_override: Optional[ResearchArtifacts] = None,
        db_session: Optional[AsyncSession] = None,
        request_id: Optional[str] = None,
    ) -> PlanPipelineArtifacts:
        if not AI_REPORTS_ENABLED:
            raise PlannerUnavailable("AI reports feature disabled by configuration")

        pipeline_start = time.perf_counter()
        prompt_tokens = _tokenize(prompt)

        schema_ms = 0
        if schema_context is None:
            schema_start = time.perf_counter()
            schema_context = await self.build_schema_context(
                account_id=account_id,
                gateway=gateway,
                prompt=prompt,
                db_session=db_session,
            )
            schema_ms = int((time.perf_counter() - schema_start) * 1000)

        research_cache_hit = False
        research_override_used = research_override is not None
        if research_override_used:
            research = research_override
            research_ms = 0
            await self._cache_research(account_id, prompt, research)
        else:
            research = await self._get_cached_research(account_id, prompt)
            if research:
                research_cache_hit = True
                research_ms = 0
            else:
                research_start = time.perf_counter()
                research = await self._run_researcher_consolidated(
                    account_id=account_id,
                    prompt=prompt,
                    schema_context=schema_context,
                )
                research_ms = int((time.perf_counter() - research_start) * 1000)
                await self._cache_research(account_id, prompt, research)
        rank_start = time.perf_counter()
        ranked_fields: List[Dict[str, Any]] = []
        index_meta: Optional[Dict[str, Any]] = None
        if db_session:
            repo = SalesforceSchemaIndexRepo(db_session)
            try:
                ranked_fields = await self._rank_fields_from_index(
                    repo=repo,
                    org_id=account_id,
                    prompt=prompt,
                )
            except Exception as exc:  # noqa: BLE001
                log.debug("Failed to rank fields from index: %s", exc)
            try:
                index_meta = await repo.status(
                    account_id,
                    ttl_hours=AI_SCHEMA_INDEX_TTL_HOURS,
                )
            except Exception as exc:  # noqa: BLE001
                log.debug("Failed to fetch schema index status: %s", exc)
        rank_ms = int((time.perf_counter() - rank_start) * 1000)
        if ranked_fields:
            research.ranked_fields = ranked_fields
            top_ranked = ranked_fields[: min(12, len(ranked_fields))]
            research.suggested_fields = self._build_suggested_fields_from_ranked(
                top_ranked
            )
            if not research.suggested_fields and ranked_fields:
                log.debug(
                    "Ranked fields present but no valid suggestions generated: %s",
                    top_ranked,
                )
            research.clarifiers = self._build_discovery_clarifiers(
                ranked_fields, prompt_tokens
            )
            research.clarifying_questions = [
                clarifier["question"] for clarifier in research.clarifiers
            ]
        if index_meta:
            research.index_status = {
                "status": index_meta.get("index_status"),
                "source": index_meta.get("index_source"),
                "stage": index_meta.get("stage"),
                "in_progress": index_meta.get("in_progress", False),
                "schema_hash": index_meta.get("schema_hash"),
            }

        planner_ms = 0
        critic_ms = 0
        critic_skipped = False
        fallback_ms = 0
        planner_fallback = False
        failure_stage: Optional[str] = None
        planner_stage = "planner"
        planner_start = time.perf_counter()
        critic_start: Optional[float] = None
        try:
            planner_draft = await self._run_planner(
                prompt=prompt,
                schema_context=schema_context,
                research=research,
                tier=tier,
                prior_warnings=prior_warnings,
                conversation=conversation,
                clarifier_answers=clarifier_answers,
            )
            planner_ms = int((time.perf_counter() - planner_start) * 1000)
            plan = self._prune_unused_joins(planner_draft.plan)
            if AI_REPORTS_SKIP_CRITIC_SIMPLE and not self._plan_requires_critic(plan):
                critic_skipped = True
                risks = []
                alternatives = []
            else:
                planner_stage = "critic"
                critic_start = time.perf_counter()
                critic_feedback = await self._run_critic(
                    prompt=prompt,
                    schema_context=schema_context,
                    research=research,
                    plan=plan,
                    tier=tier,
                )
                critic_ms = int((time.perf_counter() - critic_start) * 1000)
                risks = critic_feedback.risks
                alternatives = critic_feedback.alternatives
            clarifiers = planner_draft.clarifiers
            trace = planner_draft.trace
        except PlannerUnavailable as exc:
            planner_fallback = True
            failure_stage = planner_stage
            if planner_stage == "critic" and critic_start is not None:
                critic_ms = int((time.perf_counter() - critic_start) * 1000)
            else:
                planner_ms = int((time.perf_counter() - planner_start) * 1000)
            log.warning("Planner unavailable, returning heuristic plan: %s", exc)
            fallback_start = time.perf_counter()
            (
                plan,
                fallback_trace,
                fallback_warnings,
                fallback_risks,
                fallback_alts,
            ) = self._build_default_plan_payload(
                prompt=prompt,
                schema_context=schema_context,
                research=research,
                ranked_fields=ranked_fields,
                prompt_terms=prompt_tokens,
                org_id=account_id,
                request_id=request_id,
            )
            fallback_ms = int((time.perf_counter() - fallback_start) * 1000)
            trace = fallback_trace
            clarifiers = []
            risks = list(dict.fromkeys(fallback_warnings + fallback_risks))
            alternatives = fallback_alts

        postprocess_start = time.perf_counter()
        plan = self._ensure_defaults(
            plan=plan,
            ranked_fields=ranked_fields,
            prompt_terms=prompt_tokens,
            org_id=account_id,
            request_id=request_id,
        )
        plan = self._enforce_aggregation_rules(plan)
        plan = self._prune_unused_joins(plan)

        auto_alternatives = self._synthesize_alternatives(
            plan=plan, ranked_fields=ranked_fields
        )
        if alternatives:
            existing_titles = {alt.title for alt in alternatives}
        else:
            alternatives = []
            existing_titles = set()
        for alt in auto_alternatives:
            if alt.title not in existing_titles:
                alternatives.append(alt)
                existing_titles.add(alt.title)

        research_payload = research
        if ranked_fields:
            research_payload.ranked_fields = ranked_fields
            discovery_clarifiers = self._build_discovery_clarifiers(
                ranked_fields, prompt_tokens
            )
            research_payload.clarifiers = discovery_clarifiers
            research_payload.clarifying_questions = [
                item["question"] for item in discovery_clarifiers
            ]
        if index_meta:
            research_payload.index_status = {
                "status": index_meta.get("index_status"),
                "source": index_meta.get("index_source"),
                "stage": index_meta.get("stage"),
                "in_progress": index_meta.get("in_progress", False),
                "schema_hash": index_meta.get("schema_hash"),
            }
        postprocess_ms = int((time.perf_counter() - postprocess_start) * 1000)

        try:
            emit(
                "planner.alternatives",
                request_id=request_id,
                org_id=account_id,
                count=len(alternatives or []),
            )
        except Exception:  # noqa: BLE001
            log.debug("Failed to emit planner alternatives telemetry", exc_info=True)

        total_ms = int((time.perf_counter() - pipeline_start) * 1000)
        try:
            emit(
                "ai_reports.plan.pipeline_ms",
                request_id=request_id,
                org_id=account_id,
                tier=tier,
                schema_ms=schema_ms,
                researcher_ms=research_ms,
                researcher_cached=research_cache_hit,
                researcher_override=research_override_used,
                rank_ms=rank_ms,
                planner_ms=planner_ms,
                critic_ms=critic_ms,
                critic_skipped=critic_skipped,
                fallback_ms=fallback_ms,
                postprocess_ms=postprocess_ms,
                total_ms=total_ms,
                planner_fallback=planner_fallback,
                failure_stage=failure_stage,
                planner_model=self.select_model_for_tier(tier),
                researcher_model=AI_REPORTS_RESEARCHER_MODEL or AI_REPORTS_PLANNER_MODEL,
                critic_model=AI_REPORTS_CRITIC_MODEL
                or AI_REPORTS_RESEARCHER_MODEL
                or self.select_model_for_tier(tier),
            )
        except Exception:  # noqa: BLE001
            log.debug("Failed to emit plan pipeline telemetry", exc_info=True)

        return PlanPipelineArtifacts(
            plan=plan,
            trace=trace,
            clarifiers=clarifiers,
            risks=risks,
            alternatives=alternatives,
            research=research_payload,
        )

    async def discover(
        self,
        *,
        account_id: str,
        prompt: str,
        gateway: SalesforceGateway,
        db_session: Optional[AsyncSession] = None,
        request_id: Optional[str] = None,
    ) -> Tuple[ResearchArtifacts, SchemaPromptContext]:
        prompt_terms = _tokenize(prompt)
        schema_context = await self.build_schema_context(
            account_id=account_id,
            gateway=gateway,
            prompt=prompt,
            db_session=db_session,
        )
        research = await self._run_researcher_consolidated(
            account_id=account_id,
            prompt=prompt,
            schema_context=schema_context,
        )
        ranked_fields: List[Dict[str, Any]] = []
        index_meta: Optional[Dict[str, Any]] = None
        if db_session:
            repo = SalesforceSchemaIndexRepo(db_session)
            try:
                ranked_fields = await self._rank_fields_from_index(
                    repo=repo,
                    org_id=account_id,
                    prompt=prompt,
                )
            except Exception as exc:  # noqa: BLE001
                log.debug("Failed to generate ranked discovery fields: %s", exc)
            try:
                index_meta = await repo.status(
                    account_id,
                    ttl_hours=AI_SCHEMA_INDEX_TTL_HOURS,
                )
            except Exception as exc:  # noqa: BLE001
                log.debug("Failed to fetch schema index status: %s", exc)
        if ranked_fields:
            research.ranked_fields = ranked_fields
            top_ranked = ranked_fields[: min(12, len(ranked_fields))]
            research.suggested_fields = self._build_suggested_fields_from_ranked(
                top_ranked
            )
            if not research.suggested_fields and ranked_fields:
                log.debug(
                    "Ranked fields present but no valid suggestions generated: %s",
                    top_ranked,
                )
            research.clarifiers = self._build_discovery_clarifiers(
                ranked_fields, prompt_terms
            )
            research.clarifying_questions = [
                clarifier["question"] for clarifier in research.clarifiers
            ]
        if index_meta:
            research.index_status = {
                "status": index_meta.get("index_status"),
                "source": index_meta.get("index_source"),
                "stage": index_meta.get("stage"),
                "in_progress": index_meta.get("in_progress", False),
                "schema_hash": index_meta.get("schema_hash"),
            }
        if not research.macros:
            research.macros = self.list_macros()
        await self._cache_research(account_id, prompt, research)
        return research, schema_context

    async def _rank_fields_from_index(
        self,
        *,
        repo: SalesforceSchemaIndexRepo,
        org_id: str,
        prompt: str,
        limit: int = 20,
    ) -> List[Dict[str, Any]]:
        terms = _tokenize(prompt)
        ranked: List[Dict[str, Any]] = []
        for (
            field_ref,
            score,
            label,
            object_api,
            field_api,
            report_count,
            fill_rate,
        ) in await repo.rank_for_terms(
            org_id=org_id,
            terms=terms,
            limit=limit,
        ):
            entry = {
                "field_ref": field_ref,
                "score": score,
                "label": label,
                "object_api": object_api,
                "field_api": field_api,
                "report_usage": report_count,
                "fill_rate": fill_rate,
            }
            entry["badges"] = self._clarifier_badges(entry)
            entry["why"] = self._format_ranked_field_reason(entry)
            ranked.append(entry)
        return ranked

    def _build_suggested_fields_from_ranked(
        self, ranked_fields: Iterable[Dict[str, Any]]
    ) -> List[SuggestedField]:
        """Convert ranked field metadata into SuggestedField models with validation safeguards."""
        suggestions: List[SuggestedField] = []
        for item in ranked_fields or []:
            object_api = (
                item.get("object_api") or item.get("objectApi") or item.get("object")
            )
            field_ref = (
                item.get("field_ref") or item.get("fieldRef") or item.get("field")
            )
            if not object_api or not field_ref:
                log.debug("Skipping ranked field without object/field: %s", item)
                continue
            try:
                suggestions.append(
                    SuggestedField(
                        object=object_api,
                        field=field_ref,
                        why=item.get("why") or self._format_ranked_field_reason(item),
                        score=item.get("score"),
                        field_ref=item.get("field_ref") or item.get("fieldRef"),
                        object_api=item.get("object_api") or item.get("objectApi"),
                        report_usage=item.get("report_usage")
                        or item.get("reportUsage"),
                        fill_rate=item.get("fill_rate") or item.get("fillRate"),
                        badges=item.get("badges") or self._clarifier_badges(item),
                    )
                )
            except ValidationError as exc:
                log.debug("Skipping ranked field due to validation error: %s", exc)
        return suggestions

    def _format_ranked_field_reason(self, item: Dict[str, Any]) -> str:
        reasons: List[str] = []
        if item.get("report_usage"):
            reasons.append(f"{item['report_usage']} reports")
        fill = item.get("fill_rate")
        if isinstance(fill, (int, float)) and fill > 0:
            reasons.append(f"{int(fill * 100)}% filled")
        if item.get("score"):
            reasons.append(f"score {item['score']:.2f}")
        return " | ".join(reasons) or "High relevance"

    def _build_discovery_clarifiers(
        self,
        ranked_fields: List[Dict[str, Any]],
        prompt_terms: List[str],
    ) -> List[Dict[str, Any]]:
        clarifiers: List[Dict[str, Any]] = []
        keywords = {"amount", "arr", "carr", "net arr", "acv", "mrr", "tcv", "bookings"}
        amount_candidates = []
        for item in ranked_fields:
            label = (item.get("label") or "").lower()
            field_api = (item.get("field_api") or "").lower()
            if any(key in label for key in keywords) or any(
                key in field_api for key in keywords
            ):
                amount_candidates.append(self._clarifier_choice_dict(item))
        amount_candidates.sort(
            key=lambda choice: choice.get("score") or 0.0, reverse=True
        )
        if len(amount_candidates) >= 2:
            choices = amount_candidates[:4]
            choices.append(
                {
                    "label": "Both (two columns)",
                    "fieldRef": "__both__",
                    "badges": ["Adds ARR and Amount"],
                    "score": 0.0,
                }
            )
            clarifiers.append(
                {
                    "id": "amount_metric",
                    "question": "Which amount metric do you want?",
                    "choices": choices,
                }
            )

        prompt_tokens = {term.lower() for term in prompt_terms}
        wants_fiscal = any(
            token in {"fiscal", "fiscalperiod", "fiscalquarter", "fiscalyear"}
            for token in prompt_tokens
        )
        date_choices = [
            {
                "label": "This Fiscal Quarter",
                "value": "THIS_FISCAL_QUARTER",
                "badges": ["Custom fiscal"],
            },
            {
                "label": "Next Fiscal Quarter",
                "value": "NEXT_FISCAL_QUARTER",
                "badges": ["Custom fiscal"],
            },
            {
                "label": "Last Fiscal Quarter",
                "value": "LAST_FISCAL_QUARTER",
                "badges": ["Custom fiscal"],
            },
            {
                "label": "This Quarter",
                "value": "THIS_QUARTER",
            },
            {
                "label": "Next Quarter",
                "value": "NEXT_QUARTER",
            },
            {
                "label": "Last Quarter",
                "value": "LAST_QUARTER",
            },
            {
                "label": "Last 90 Days",
                "value": "LAST_N_DAYS:90",
            },
        ]
        if not wants_fiscal:
            date_choices = [
                choice
                for choice in date_choices
                if not choice["value"].endswith("FISCAL_QUARTER")
            ]
        preferred_literal = self._default_close_date_literal(list(prompt_tokens))
        preferred_choices = [
            choice for choice in date_choices if choice.get("value") == preferred_literal
        ]
        if preferred_choices:
            date_choices = preferred_choices + [
                choice
                for choice in date_choices
                if choice.get("value") != preferred_literal
            ]
        clarifiers.append(
            {
                "id": "date_window",
                "question": "Which date window?",
                "choices": date_choices,
            }
        )
        return clarifiers

    def _default_close_date_literal(self, prompt_terms: List[str]) -> str:
        tokens = {term.lower() for term in (prompt_terms or [])}
        wants_fiscal = any(
            token in {"fiscal", "fiscalperiod", "fiscalquarter", "fiscalyear"}
            for token in tokens
        )
        mentions_quarter = any(
            token in {"quarter", "quarters", "qtr", "qtrs", "fiscalquarter"}
            for token in tokens
        )
        wants_next = any(token in {"next", "upcoming"} for token in tokens)
        wants_last = any(token in {"last", "previous", "prior"} for token in tokens)

        if wants_fiscal and mentions_quarter:
            if wants_next:
                return "NEXT_FISCAL_QUARTER"
            if wants_last:
                return "LAST_FISCAL_QUARTER"
            return "THIS_FISCAL_QUARTER"
        if mentions_quarter:
            if wants_next:
                return "NEXT_QUARTER"
            if wants_last:
                return "LAST_QUARTER"
            return "THIS_QUARTER"
        return "THIS_FISCAL_QUARTER" if wants_fiscal else "THIS_QUARTER"

    def _clarifier_choice_dict(self, item: Dict[str, Any]) -> Dict[str, Any]:
        return {
            "label": item.get("label") or item.get("field_api"),
            "fieldRef": item.get("field_ref"),
            "score": item.get("score") or 0.0,
            "badges": self._clarifier_badges(item),
            "reportUsage": item.get("report_usage") or 0,
            "fillRate": item.get("fill_rate") or 0.0,
        }

    def _clarifier_badges(self, item: Dict[str, Any]) -> List[str]:
        badges: List[str] = []
        if item.get("report_usage"):
            badges.append(f"{item['report_usage']} reports")
        fill = item.get("fill_rate")
        if isinstance(fill, (int, float)) and fill > 0:
            badges.append(f"{int(fill * 100)}% filled")
        field_api = (item.get("field_api") or "").lower()
        if field_api == "amount":
            badges.append("Standard")
        elif field_api.endswith("__c"):
            badges.append("Custom")
        return badges

    def _ensure_defaults(
        self,
        *,
        plan: SalesforcePlan,
        ranked_fields: List[Dict[str, Any]],
        prompt_terms: List[str],
        org_id: Optional[str] = None,
        request_id: Optional[str] = None,
    ) -> SalesforcePlan:
        updated = plan.model_copy(deep=True)
        primary = updated.objects[0] if updated.objects else "Opportunity"
        existing_fields = list(updated.fields or [])
        existing_set = set(existing_fields)

        best_amount = self._best_amount_field(ranked_fields, allow_standard=True)
        amount_choice: Optional[str] = None
        if best_amount and best_amount not in existing_set:
            existing_fields.append(best_amount)
            existing_set.add(best_amount)
            amount_choice = best_amount
        elif f"{primary}.Amount" not in existing_set:
            default_amount = f"{primary}.Amount"
            existing_fields.append(default_amount)
            existing_set.add(default_amount)
            amount_choice = default_amount

        is_pipeline = any(
            (flt.field or "").lower().endswith("isclosed")
            for flt in (updated.filters or [])
        )
        added_owner = False
        added_stage = False
        if is_pipeline:
            owner_field = f"{primary}.Owner.Name"
            stage_field = f"{primary}.StageName"
            if stage_field not in existing_set:
                existing_fields.append(stage_field)
                existing_set.add(stage_field)
                added_stage = True
            if owner_field not in existing_set:
                existing_fields.append(owner_field)
                existing_set.add(owner_field)
                added_owner = True

        has_close_filter = any(
            (flt.field or "").lower().endswith("closedate")
            for flt in (updated.filters or [])
        )
        literal_choice: Optional[str] = None
        if not has_close_filter and primary.lower() == "opportunity":
            literal = self._default_close_date_literal(prompt_terms)
            literal_choice = literal
            updated.filters.append(
                PlanFilter(
                    field=f"{primary}.CloseDate", op="DATE_LITERAL", value=literal
                )
            )
        else:
            literal_choice = "existing"

        updated.fields = list(dict.fromkeys(existing_fields))
        literal_bucket = None
        if literal_choice == "THIS_FISCAL_QUARTER":
            literal_bucket = "fiscal"
        elif literal_choice == "THIS_QUARTER":
            literal_bucket = "calendar"
        elif literal_choice:
            literal_bucket = literal_choice

        if org_id:
            try:
                emit(
                    "planner.defaults",
                    request_id=request_id,
                    org_id=org_id,
                    amount=amount_choice,
                    date_literal=literal_bucket,
                    pipeline=is_pipeline,
                    added_owner=added_owner,
                    added_stage=added_stage,
                )
            except Exception:  # noqa: BLE001
                log.debug("Failed to emit planner defaults telemetry", exc_info=True)
        return updated

    def _best_amount_field(
        self,
        ranked_fields: List[Dict[str, Any]],
        *,
        allow_standard: bool,
    ) -> Optional[str]:
        keywords = {"amount", "arr", "carr", "net arr", "acv", "mrr", "tcv", "bookings"}
        candidates: List[Dict[str, Any]] = []
        for item in ranked_fields:
            label = (item.get("label") or "").lower()
            field_api = (item.get("field_api") or "").lower()
            if not allow_standard and field_api == "amount":
                continue
            if any(key in label for key in keywords) or any(
                key in field_api for key in keywords
            ):
                candidates.append(item)
        if not candidates:
            return None
        top = max(candidates, key=lambda x: x.get("score") or 0.0)
        return top.get("field_ref")

    def _synthesize_alternatives(
        self,
        *,
        plan: SalesforcePlan,
        ranked_fields: List[Dict[str, Any]],
    ) -> List[PlanAlternative]:
        alternatives: List[PlanAlternative] = []
        plan_fields = plan.fields or []
        arr_field = self._best_amount_field(ranked_fields, allow_standard=False)
        if arr_field and any(field.endswith(".Amount") for field in plan_fields):
            try:
                idx = next(
                    i
                    for i, field in enumerate(plan_fields)
                    if field.endswith(".Amount")
                )
            except StopIteration:
                idx = None
            if idx is not None:
                patch = PlanPatch(
                    ops=[
                        PlanPatchOp(
                            op="replace", path=f"/fields/{idx}", value=arr_field
                        )
                    ]
                )
                alternatives.append(
                    PlanAlternative(
                        title="Use ARR metric",
                        description="Swap the default Amount column with the highest-scoring ARR field.",
                        patch=_ensure_patch_summary(patch),
                    )
                )
        primary = plan.objects[0] if plan.objects else "Opportunity"
        forecast_field = f"{primary}.ForecastCategoryName"
        if forecast_field not in plan_fields:
            patch = PlanPatch(
                ops=[PlanPatchOp(op="add", path="/fields/-", value=forecast_field)]
            )
            alternatives.append(
                PlanAlternative(
                    title="Add Forecast category",
                    description="Include ForecastCategoryName to align with pipeline reviews.",
                    patch=_ensure_patch_summary(patch),
                )
            )
        return alternatives

    def apply_patch(self, plan: SalesforcePlan, patch: PlanPatch) -> SalesforcePlan:
        if not patch or not patch.ops:
            return plan
        payload = plan.model_dump(by_alias=True)
        mutated = _apply_patch_operations(payload, patch)
        try:
            updated_plan = SalesforcePlan.model_validate(mutated)
        except ValidationError as exc:
            raise PlanValidationError(f"Patch produced invalid plan: {exc}") from exc
        return updated_plan

    def suggest_autofix(
        self,
        *,
        plan: SalesforcePlan,
        error_message: str,
    ) -> List[PlanAlternative]:
        message = error_message or ""
        suggestions: List[PlanAlternative] = []
        contact_relation = re.search(
            r"No such relation ([A-Za-z0-9_]+) on ([A-Za-z0-9_]+)", message
        )
        if contact_relation:
            relation = contact_relation.group(1)
            base_object = contact_relation.group(2)
            if (
                relation.lower() == "contact"
                and plan.objects
                and base_object.lower() == plan.objects[0].lower()
            ):
                for idx, field in enumerate(plan.fields or []):
                    if "Contact." in field:
                        suggestions.append(
                            PlanAlternative(
                                title="Use OpportunityContactRoles.Contact.Email",
                                description="Switch to the primary contact via OpportunityContactRoles.",
                                patch=_ensure_patch_summary(
                                    PlanPatch(
                                        ops=[
                                            PlanPatchOp(
                                                op="add",
                                                path="/joins/-",
                                                value={
                                                    "from": plan.objects[0],
                                                    "to": "OpportunityContactRole",
                                                    "via": "Id",
                                                    "relationship": "OpportunityContactRoles",
                                                },
                                            ),
                                            PlanPatchOp(
                                                op="replace",
                                                path=f"/fields/{idx}",
                                                value="OpportunityContactRoles.Contact.Email",
                                            ),
                                            PlanPatchOp(
                                                op="add",
                                                path="/filters/-",
                                                value={
                                                    "field": "OpportunityContactRoles.IsPrimary",
                                                    "op": "=",
                                                    "value": True,
                                                },
                                            ),
                                        ]
                                    )
                                ),
                            )
                        )
                        suggestions.append(
                            PlanAlternative(
                                title="Drop Contact.Email field",
                                description="Remove the unavailable contact relationship from the plan.",
                                patch=_ensure_patch_summary(
                                    PlanPatch(
                                        ops=[
                                            PlanPatchOp(
                                                op="remove", path=f"/fields/{idx}"
                                            )
                                        ]
                                    )
                                ),
                            )
                        )
                        break
        if not suggestions:
            # Generic fallback: suggest adding limit or removing last field referenced in error
            if plan.fields:
                suggestions.append(
                    PlanAlternative(
                        title="Remove last field",
                        description="Drop the most recent field to satisfy Salesforce validation.",
                        patch=_ensure_patch_summary(
                            PlanPatch(
                                ops=[
                                    PlanPatchOp(
                                        op="remove",
                                        path=f"/fields/{len(plan.fields) - 1}",
                                    )
                                ]
                            )
                        ),
                    )
                )
        return suggestions[:2]

    async def peek_schema_field(
        self,
        *,
        account_id: str,
        gateway: SalesforceGateway,
        object_name: str,
        field_name: str,
        sample_limit: int = 3,
    ) -> Dict[str, Any]:
        schema_obj = await self.schema_index._fetch_object(
            account_id, object_name, gateway
        )
        schema_field = schema_obj.fields.get(field_name.lower())
        if not schema_field:
            raise PlanValidationError(f"Unknown field {object_name}.{field_name}")

        result: Dict[str, Any] = {
            "object": schema_obj.name,
            "field": schema_field.name,
            "label": schema_field.label,
            "type": schema_field.data_type,
            "helpText": schema_field.help_text,
            "relationshipName": schema_field.relationship_name,
            "referenceTo": schema_field.reference_to,
            "accessible": schema_field.accessible,
            "picklistValues": schema_field.picklist_values[:20],
            "length": schema_field.length,
            "sampleValues": [],
        }

        limit = max(1, min(sample_limit, 5))
        try:
            soql = f"SELECT {schema_field.name} FROM {schema_obj.name} WHERE {schema_field.name} != null LIMIT {limit}"
            query_result = await gateway.soql(
                soql,
                admin_override=False,
                user=f"ai-reports@{account_id}",
                ip="0.0.0.0",
                user_is_admin=False,
                plan_allowlist=None,
            )
            for record in query_result.get("records", []) or []:
                cell = record.get(schema_field.name)
                if isinstance(cell, dict):
                    value = (
                        cell.get("Name")
                        or cell.get("Id")
                        or json.dumps(cell, ensure_ascii=False)
                    )
                else:
                    value = cell
                if value is not None:
                    result["sampleValues"].append(value)
                if len(result["sampleValues"]) >= limit:
                    break
        except Exception as exc:  # noqa: BLE001
            log.debug(
                "Failed to fetch sample values for %s.%s: %s",
                object_name,
                field_name,
                exc,
            )
        return result

    @staticmethod
    def _split_field_reference(reference: str, primary: str) -> Tuple[str, str]:
        if not reference:
            return primary, ""
        parts = reference.split(".")
        if len(parts) > 1 and parts[0][:1].isupper():
            return parts[0], parts[-1]
        return primary, parts[-1]

    def _enforce_caps(self, plan: SalesforcePlan) -> None:
        unique_objects = {obj for obj in plan.objects if obj}
        if len(unique_objects) > AI_REPORTS_MAX_OBJECTS:
            raise PlanValidationError(
                f"Too many objects requested; max {AI_REPORTS_MAX_OBJECTS}"
            )
        total_fields = len(plan.fields or [])
        if total_fields > AI_REPORTS_MAX_FIELDS:
            raise PlanValidationError(
                f"Too many fields requested; max {AI_REPORTS_MAX_FIELDS}"
            )

    def _translate_filters(self, plan: SalesforcePlan) -> None:
        if not plan.filters:
            return
        from fmatch.saas.utils.soql_dates import to_literal

        for flt in plan.filters:
            if flt.op.upper() == "DATE_LITERAL":
                literal = to_literal(str(flt.value or "")) or str(flt.value or "")
                flt.value = literal

    def _prune_inaccessible_fields(
        self,
        plan: SalesforcePlan,
        schema_context: SchemaPromptContext,
    ) -> Tuple[List[str], Dict[str, set[str]]]:
        allow_override = {
            str(obj).lower(): {str(field).lower() for field in (fields or [])}
            for obj, fields in (AI_REPORTS_ALLOW_FIELDS or {}).items()
            if isinstance(fields, (list, tuple, set))
        }

        accessible_map: Dict[str, set[str]] = {}
        canonical_names: Dict[str, str] = {}
        for obj_lower, schema_obj in (schema_context.objects_by_name or {}).items():
            accessible_fields = {
                field.name.lower()
                for field in schema_obj.fields.values()
                if field.queryable and field.accessible
            }
            override = allow_override.get(obj_lower)
            if override:
                if accessible_fields:
                    accessible_fields &= override
                else:
                    accessible_fields = set(override)
            accessible_map[obj_lower] = accessible_fields
            canonical_names[obj_lower] = schema_obj.name

        removed: List[str] = []
        primary = plan.objects[0] if plan.objects else "Lead"
        primary_lower = primary.lower()

        agg_pattern = re.compile(
            r"(?i)^(?:sum|avg|count|min|max)\((.+)\)(?:\s+as\s+\w+)?$"
        )

        def _strip_aggregate(expr: str) -> str:
            match = agg_pattern.match(expr or "")
            if match:
                return match.group(1).strip()
            return expr or ""

        # Build set of objects the planner explicitly included (either in objects
        # list or as join targets). These should pass accessibility checks and
        # let join validation handle whether the relationship is valid.
        plan_objects = {obj.lower() for obj in (plan.objects or [])}
        for join in plan.joins or []:
            plan_objects.add(join.to_object.lower())
            if join.relationship:
                plan_objects.add(join.relationship.lower())

        def is_accessible(obj_name: str, field_name: str) -> bool:
            obj_lower = obj_name.lower()
            allowed = accessible_map.get(obj_lower)
            if allowed is not None and allowed:
                # Object is in schema context - check field accessibility
                return field_name.lower() in allowed
            # Object NOT in schema context - allow if planner explicitly included it
            # This defers validation to _prune_unjoined_fields which checks joins
            if obj_lower in plan_objects:
                return True
            return False

        def retarget_object(
            obj_name: str, field_name: str, original_ref: str
        ) -> Optional[str]:
            """
            Try to retarget fields that landed on the primary object (whether or not
            the LLM added the prefix) to whichever object actually exposes the column.
            This covers prompts like "CloseDate" or even "Account.CloseDate" when the
            query really needs Opportunity.
            """
            if obj_name.lower() != primary_lower:
                return None
            field_lower = field_name.lower()
            candidates = [
                canonical_names.get(obj_lower, obj_lower)
                for obj_lower, allowed in accessible_map.items()
                if field_lower in allowed
            ]
            if len(candidates) == 1:
                return candidates[0]
            return None

        kept_fields: List[str] = []
        for field in plan.fields or []:
            clean_field = _strip_aggregate(field)
            obj_name, base_field = self._split_field_reference(clean_field, primary)
            target_obj = obj_name
            if not is_accessible(obj_name, base_field):
                retargeted = retarget_object(obj_name, base_field, field)
                if retargeted:
                    target_obj = retargeted
            if is_accessible(target_obj, base_field):
                canonical = canonical_names.get(target_obj.lower(), target_obj)
                if field.count(".") == 0 or target_obj != obj_name:
                    kept_fields.append(f"{canonical}.{base_field}")
                else:
                    kept_fields.append(field)
            else:
                removed.append(f"{obj_name}.{base_field}")
        plan.fields = kept_fields

        filtered_filters: List[PlanFilter] = []
        for flt in plan.filters or []:
            base_expr = _strip_aggregate(flt.field)
            obj_name, base_field = self._split_field_reference(base_expr, primary)
            target_obj = obj_name
            if not is_accessible(obj_name, base_field):
                retargeted = retarget_object(obj_name, base_field, flt.field)
                if retargeted:
                    target_obj = retargeted
                    canonical = canonical_names.get(target_obj.lower(), target_obj)
                    flt.field = f"{canonical}.{base_field}"
            if is_accessible(target_obj, base_field):
                filtered_filters.append(flt)
            else:
                removed.append(f"Filter on {obj_name}.{base_field} (access denied)")
        plan.filters = filtered_filters

        filtered_orders: List[PlanOrder] = []
        for order in plan.order_by or []:
            clean_field = _strip_aggregate(order.field)
            obj_name, base_field = self._split_field_reference(clean_field, primary)
            target_obj = obj_name
            if not is_accessible(obj_name, base_field):
                retargeted = retarget_object(obj_name, base_field, order.field)
                if retargeted:
                    target_obj = retargeted
                    canonical = canonical_names.get(target_obj.lower(), target_obj)
                    order.field = f"{canonical}.{base_field}"
            if is_accessible(target_obj, base_field):
                filtered_orders.append(order)
            else:
                removed.append(f"{obj_name}.{base_field} (order)")
        plan.order_by = filtered_orders

        filtered_group: List[str] = []
        for group_field in plan.group_by or []:
            clean_field = _strip_aggregate(group_field)
            obj_name, base_field = self._split_field_reference(clean_field, primary)
            target_obj = obj_name
            if not is_accessible(obj_name, base_field):
                retargeted = retarget_object(obj_name, base_field, group_field)
                if retargeted:
                    target_obj = retargeted
                    canonical = canonical_names.get(target_obj.lower(), target_obj)
                    filtered_group.append(f"{canonical}.{base_field}")
                    continue
            if is_accessible(target_obj, base_field):
                filtered_group.append(group_field)
            else:
                removed.append(f"{obj_name}.{base_field} (group)")
        plan.group_by = filtered_group

        return removed, accessible_map

    def _prune_unjoined_fields(
        self, plan: SalesforcePlan
    ) -> Tuple[List[str], Dict[str, List[str]]]:
        """Remove fields that reference objects without corresponding joins.

        When a field like 'Campaign.AmountAllOpportunities' is in the plan
        but there's no join from the primary object to Campaign, the SOQL
        will fail. This method identifies and removes such fields.

        Returns:
            Tuple of:
            - List of removed field references
            - Dict mapping missing object names to list of fields that need them
        """
        if not plan.objects:
            return [], {}

        primary = plan.objects[0].lower()
        primary_title = plan.objects[0]  # Keep original case for suggestions

        # Standard polymorphic relationships available on ALL objects
        # These are always traversable without explicit joins
        # NOTE: Do NOT include Account/Contact here - they're NOT universal!
        # Lead has no Account lookup (only Company text field), so allowing
        # "account" causes invalid SOQL like "SELECT Account.Name FROM Lead"
        standard_relationships = {
            "owner", "createdby", "lastmodifiedby",
            # Common relationship names that resolve to User
            "user", "manager",
        }

        # Common Salesforce parent-child lookup relationships
        # Format: {(from_object, to_object): (via_field, relationship_name)}
        # These allow auto-generation of joins when the planner forgets
        common_lookups: Dict[Tuple[str, str], Tuple[str, str]] = {
            # Opportunity relationships
            ("opportunity", "account"): ("AccountId", "Account"),
            ("opportunity", "campaign"): ("CampaignId", "Campaign"),
            ("opportunity", "contact"): ("ContactId", "Contact"),
            ("opportunity", "pricebook2"): ("Pricebook2Id", "Pricebook2"),
            # OpportunityLineItem relationships
            ("opportunitylineitem", "opportunity"): ("OpportunityId", "Opportunity"),
            ("opportunitylineitem", "product2"): ("Product2Id", "Product2"),
            ("opportunitylineitem", "pricebookentry"): ("PricebookEntryId", "PricebookEntry"),
            # Contact relationships
            ("contact", "account"): ("AccountId", "Account"),
            ("contact", "campaign"): ("CampaignId", "Campaign"),
            # Case relationships
            ("case", "account"): ("AccountId", "Account"),
            ("case", "contact"): ("ContactId", "Contact"),
            # Task/Event relationships
            ("task", "account"): ("AccountId", "Account"),
            ("task", "contact"): ("WhoId", "Who"),
            ("event", "account"): ("AccountId", "Account"),
            ("event", "contact"): ("WhoId", "Who"),
            # CampaignMember relationships
            ("campaignmember", "campaign"): ("CampaignId", "Campaign"),
            ("campaignmember", "contact"): ("ContactId", "Contact"),
            ("campaignmember", "lead"): ("LeadId", "Lead"),
        }

        def try_auto_join(from_obj: str, to_obj: str) -> Optional[PlanJoin]:
            """Try to auto-generate a join using common Salesforce relationships."""
            key = (from_obj.lower(), to_obj.lower())
            if key in common_lookups:
                via, relationship = common_lookups[key]
                return PlanJoin(
                    from_object=from_obj,
                    to_object=to_obj,
                    via=via,
                    relationship=relationship,
                )
            return None

        # Build set of reachable objects STRICTLY from:
        # 1. Primary object only
        # 2. Explicit joins defined in the plan
        # 3. Standard polymorphic relationships
        # NOTE: Do NOT blindly trust plan.objects[1:] - this causes Object Drift
        # when the LLM hallucinates objects that have no actual join path
        reachable = {primary}
        reachable.update(standard_relationships)

        # Add objects reachable via explicit joins only
        # Both from_object and to_object are reachable since the join defines a path
        for join in plan.joins or []:
            reachable.add(join.from_object.lower())
            reachable.add(join.to_object.lower())
            # Also add the relationship path as reachable (e.g., "Account" from "Account.Name")
            if join.relationship:
                reachable.add(join.relationship.lower())

        agg_pattern = re.compile(
            r"(?i)^(?:sum|avg|count|min|max)\((.+)\)(?:\s+as\s+\w+)?$"
        )

        def extract_object(field_ref: str) -> Tuple[Optional[str], Optional[str]]:
            """Extract the object name from a field reference.

            Returns (lowercase_name, original_case_name) tuple.
            """
            # Strip aggregate wrapper
            match = agg_pattern.match(field_ref)
            if match:
                field_ref = match.group(1).strip()
            if "." not in field_ref:
                return None, None  # Unqualified field, assumed primary
            parts = field_ref.split(".")
            # Handle traversals like Owner.Name -> Owner
            return parts[0].lower(), parts[0]

        # ─── AUTO-JOIN GENERATION ───
        # Collect all objects referenced in the plan that aren't yet reachable
        # and try to auto-generate joins using common Salesforce relationships
        referenced_objects: Set[str] = set()
        for field in plan.fields or []:
            obj_lower, _ = extract_object(field)
            if obj_lower:
                referenced_objects.add(obj_lower)
        for flt in plan.filters or []:
            obj_lower, _ = extract_object(flt.field)
            if obj_lower:
                referenced_objects.add(obj_lower)
        for order in plan.order_by or []:
            obj_lower, _ = extract_object(order.field)
            if obj_lower:
                referenced_objects.add(obj_lower)
        for group in plan.group_by or []:
            obj_lower, _ = extract_object(group)
            if obj_lower:
                referenced_objects.add(obj_lower)

        # Try to auto-generate joins for unreachable objects
        auto_generated_joins: List[PlanJoin] = []
        for ref_obj in referenced_objects:
            if ref_obj not in reachable:
                # Try direct join from primary
                auto_join = try_auto_join(primary_title, ref_obj)
                if auto_join:
                    auto_generated_joins.append(auto_join)
                    reachable.add(ref_obj)
                    if auto_join.relationship:
                        reachable.add(auto_join.relationship.lower())
                    log.info(
                        "Auto-generated join: %s -> %s (via %s)",
                        primary_title,
                        ref_obj,
                        auto_join.via,
                    )

        # Add auto-generated joins to the plan
        if auto_generated_joins:
            if plan.joins is None:
                plan.joins = []
            plan.joins.extend(auto_generated_joins)

        removed: List[str] = []
        missing_objects: Dict[str, List[str]] = {}  # object_name -> [fields]

        # Filter fields
        kept_fields: List[str] = []
        for field in plan.fields or []:
            obj_lower, obj_original = extract_object(field)
            if obj_lower is None or obj_lower in reachable:
                kept_fields.append(field)
            else:
                removed.append(field)
                # Track which object is missing for join suggestions
                if obj_original:
                    missing_objects.setdefault(obj_original, []).append(field)
                log.info(
                    "Dropping field '%s': object '%s' has no join path from '%s'",
                    field,
                    obj_lower,
                    primary,
                )
        plan.fields = kept_fields

        # Filter filters
        kept_filters: List[PlanFilter] = []
        for flt in plan.filters or []:
            obj_lower, obj_original = extract_object(flt.field)
            if obj_lower is None or obj_lower in reachable:
                kept_filters.append(flt)
            else:
                removed.append(flt.field)
                if obj_original:
                    missing_objects.setdefault(obj_original, []).append(flt.field)
                log.info(
                    "Dropping filter on '%s': object '%s' has no join path",
                    flt.field,
                    obj_lower,
                )
        plan.filters = kept_filters

        # Filter order_by
        kept_orders: List[PlanOrder] = []
        for order in plan.order_by or []:
            obj_lower, obj_original = extract_object(order.field)
            if obj_lower is None or obj_lower in reachable:
                kept_orders.append(order)
            else:
                removed.append(order.field)
                if obj_original:
                    missing_objects.setdefault(obj_original, []).append(order.field)
        plan.order_by = kept_orders

        # Filter group_by
        kept_groups: List[str] = []
        for group in plan.group_by or []:
            obj_lower, obj_original = extract_object(group)
            if obj_lower is None or obj_lower in reachable:
                kept_groups.append(group)
            else:
                removed.append(group)
                if obj_original:
                    missing_objects.setdefault(obj_original, []).append(group)
        plan.group_by = kept_groups

        return removed, missing_objects

    def validate_plan(
        self,
        *,
        plan: SalesforcePlan,
        schema_context: SchemaPromptContext,
        row_cap: int,
    ) -> Tuple[List[str], List[PlanAlternative]]:
        warnings: List[str] = []

        if not plan.objects:
            raise PlanValidationError("Plan must include at least one object")

        self._enforce_caps(plan)

        primary_object = plan.objects[0]
        log.debug(
            "AI Reports plan (org=%s): objects=%s joins=%s fields=%s filters=%s orders=%s",
            schema_context.prompt_objects[0].name if schema_context.prompt_objects else "unknown",
            plan.objects,
            [join.model_dump() for join in (plan.joins or [])],
            plan.fields,
            [flt.model_dump() for flt in (plan.filters or [])],
            [order.model_dump() for order in (plan.order_by or [])],
        )

        # ─── CHILD-OBJECT AGGREGATION DETECTION ───
        # Detect when the user wants to aggregate on child object fields
        # and auto-reorient the query to start from the child object
        child_object_mappings = {
            # (parent, child): new_primary
            ("opportunity", "opportunitylineitem"): "OpportunityLineItem",
            ("opportunity", "product2"): "OpportunityLineItem",  # Via LineItem
            ("account", "contact"): "Contact",
            ("account", "opportunity"): "Opportunity",
            ("campaign", "campaignmember"): "CampaignMember",
        }

        agg_pattern = re.compile(r"(?i)^(?:sum|avg|count|min|max)\((.+)\)")

        def needs_child_reorientation() -> Optional[str]:
            """Check if query should be reoriented to a child object."""
            primary_lower = primary_object.lower()
            for field in plan.fields or []:
                # Check aggregations on child object fields
                agg_match = agg_pattern.match(field)
                if agg_match:
                    inner_field = agg_match.group(1).strip()
                    if "." in inner_field:
                        obj = inner_field.split(".")[0].lower()
                        key = (primary_lower, obj)
                        if key in child_object_mappings:
                            return child_object_mappings[key]
            # Check group_by on child object fields
            for group in plan.group_by or []:
                if "." in group:
                    obj = group.split(".")[0].lower()
                    key = (primary_lower, obj)
                    if key in child_object_mappings:
                        return child_object_mappings[key]
            return None

        suggested_reorient = needs_child_reorientation()
        if suggested_reorient:
            log.info(
                "Query may benefit from reorientation: %s -> %s",
                primary_object,
                suggested_reorient,
            )
            # Note: We don't auto-reorient here, but the reviewer will suggest it

        # ─── JOIN VALIDATION ───
        # Validate that join 'via' fields exist on the 'from_object' in the schema.
        # This prevents invalid joins like Lead→Account (Lead has no AccountId).
        if plan.joins and schema_context.objects_by_name:
            valid_joins: List[PlanJoin] = []
            removed_joins: List[str] = []

            for join in plan.joins:
                from_obj_lower = join.from_object.lower()
                schema_obj = schema_context.objects_by_name.get(from_obj_lower)

                if schema_obj is None:
                    # Object not in schema - can't validate, keep the join
                    valid_joins.append(join)
                    continue

                # Check if 'via' field exists on the from_object
                via_field_lower = join.via.lower()
                field_exists = any(
                    f.name.lower() == via_field_lower
                    for f in schema_obj.fields.values()
                )

                if field_exists:
                    valid_joins.append(join)
                else:
                    removed_joins.append(
                        f"{join.from_object}→{join.to_object} (via {join.via})"
                    )
                    log.info(
                        "Removing invalid join %s→%s: field '%s' does not exist on %s",
                        join.from_object,
                        join.to_object,
                        join.via,
                        join.from_object,
                    )

            if removed_joins:
                plan.joins = valid_joins
                warnings.append(
                    "Removed invalid joins: " + ", ".join(removed_joins)
                )

        available_objects = schema_context.objects_by_name or {}
        primary_lower = primary_object.lower()
        primary_schema = available_objects.get(primary_lower)
        if primary_schema is None:
            raise PlanValidationError(
                f"Primary object '{primary_object}' is not present in schema context"
            )

        relationship_aliases: Set[str] = {
            "owner",
            "createdby",
            "lastmodifiedby",
            "user",
            "manager",
            "who",
            "what",
        }
        relationship_targets: Dict[str, str] = {}

        for rel_name, targets in (primary_schema.relationships or {}).items():
            rel_lower = str(rel_name or "").strip().lower()
            if not rel_lower:
                continue
            relationship_aliases.add(rel_lower)
            for target in targets or []:
                target_lower = str(target or "").strip().lower()
                if target_lower in available_objects:
                    relationship_targets.setdefault(rel_lower, target_lower)
                    break

        for schema_field in primary_schema.fields.values():
            rel_name = str(schema_field.relationship_name or "").strip().lower()
            if not rel_name:
                continue
            relationship_aliases.add(rel_name)
            for target in schema_field.reference_to or []:
                target_lower = str(target or "").strip().lower()
                if target_lower in available_objects:
                    relationship_targets.setdefault(rel_name, target_lower)
                    break

        for join in plan.joins or []:
            if join.relationship:
                rel_lower = join.relationship.lower()
                relationship_aliases.add(rel_lower)
                to_lower = join.to_object.lower()
                if to_lower in available_objects:
                    relationship_targets.setdefault(rel_lower, to_lower)

        def ensure_field_known(field_path: str) -> None:
            reference = _unwrap_field_reference(field_path)
            if not reference:
                return
            parts = [part.strip() for part in reference.split(".") if part.strip()]
            if not parts:
                return

            target_object_lower = primary_lower
            if len(parts) > 1:
                prefix = parts[0]
                prefix_lower = prefix.lower()
                if prefix_lower in available_objects:
                    target_object_lower = prefix_lower
                elif prefix_lower in relationship_aliases:
                    rel_target = relationship_targets.get(prefix_lower)
                    if rel_target is None:
                        # Relationship alias is known, but the target object is ambiguous.
                        # Keep runtime Salesforce validation as the source of truth here.
                        return
                    target_object_lower = rel_target
                else:
                    raise PlanValidationError(
                        f"Field '{field_path}' references unknown object or relationship '{prefix}'"
                    )

            target_schema = available_objects.get(target_object_lower)
            if target_schema is None:
                raise PlanValidationError(
                    f"Field '{field_path}' references unknown object '{target_object_lower}'"
                )

            field_name = parts[-1]
            if field_name.lower() not in target_schema.fields:
                raise PlanValidationError(
                    f"Unknown field '{field_name}' on object {target_schema.name}"
                )

        for selected_field in plan.fields:
            ensure_field_known(selected_field)
        for filt in plan.filters:
            ensure_field_known(filt.field)
        for order_clause in plan.order_by:
            ensure_field_known(order_clause.field)
        for grouping in plan.group_by:
            ensure_field_known(grouping)

        removed_fields, accessible_map = self._prune_inaccessible_fields(
            plan, schema_context
        )
        if removed_fields:
            warnings.append(
                "Removed fields without access: "
                + ", ".join(sorted(set(removed_fields)))
            )

        # Remove fields referencing objects without join paths
        unjoined_fields, missing_objects = self._prune_unjoined_fields(plan)
        if unjoined_fields:
            warnings.append(
                "Removed fields with missing joins: "
                + ", ".join(sorted(set(unjoined_fields)))
            )

        # Generate auto-fix suggestions for missing joins
        join_suggestions: List[PlanAlternative] = []
        primary = plan.objects[0] if plan.objects else "Opportunity"

        # If query needs reorientation to child object, add as first suggestion
        if suggested_reorient:
            join_suggestions.append(
                PlanAlternative(
                    title=f"Reorient to {suggested_reorient}",
                    description=(
                        f"Start query from {suggested_reorient} instead of {primary_object}. "
                        f"This allows direct aggregation of child object fields in SOQL."
                    ),
                    patch=_ensure_patch_summary(
                        PlanPatch(
                            ops=[
                                PlanPatchOp(
                                    op="replace",
                                    path="/objects/0",
                                    value=suggested_reorient,
                                ),
                            ]
                        )
                    ),
                )
            )

        for obj_name, fields in missing_objects.items():
            field_list = ", ".join(fields[:3])
            if len(fields) > 3:
                field_list += f" (+{len(fields) - 3} more)"
            join_suggestions.append(
                PlanAlternative(
                    title=f"Add {obj_name} join",
                    description=f"Join {obj_name} to include: {field_list}",
                    patch=_ensure_patch_summary(
                        PlanPatch(
                            ops=[
                                PlanPatchOp(
                                    op="add",
                                    path="/joins/-",
                                    value={
                                        "from": primary,
                                        "to": obj_name,
                                        "via": f"{obj_name}Id",
                                        "relationship": obj_name,
                                    },
                                ),
                            ]
                            + [
                                PlanPatchOp(op="add", path="/fields/-", value=fld)
                                for fld in fields
                            ]
                        )
                    ),
                )
            )

        if not plan.fields:
            fallback = None
            for obj_lower, schema_obj in (schema_context.objects_by_name or {}).items():
                for schema_field in schema_obj.fields.values():
                    if schema_field.queryable and schema_field.accessible:
                        fallback = f"{schema_obj.name}.{schema_field.name}"
                        break
                if fallback:
                    break
            if fallback:
                plan.fields.append(fallback)
                removed_retry, _ = self._prune_inaccessible_fields(plan, schema_context)
                if removed_retry:
                    warnings.append(
                        f"Fallback field '{fallback}' not accessible; plan may need adjustment"
                    )
            else:
                warnings.append(
                    "No accessible fields available; adjust plan or permissions."
                )

        self._translate_filters(plan)

        if plan.limit is None:
            warnings.append(f"No limit specified; defaulting to {row_cap}")
            plan.limit = row_cap

        if plan.limit <= 0:
            plan.limit = row_cap
            warnings.append(f"Adjusted invalid limit to {row_cap}")

        if plan.limit > row_cap:
            warnings.append(
                f"Limit reduced from {plan.limit} to {row_cap} for tier constraints"
            )
            plan.limit = row_cap

        # Multi-currency validation: detect Amount aggregations without CurrencyIsoCode
        currency_suggestions = self._check_multicurrency_aggregation(
            plan, schema_context, primary_object
        )
        join_suggestions.extend(currency_suggestions)

        return warnings, join_suggestions

    def _check_multicurrency_aggregation(
        self,
        plan: SalesforcePlan,
        schema_context: SchemaPromptContext,
        primary_object: str,
    ) -> List[PlanAlternative]:
        """Detect Amount aggregations in multi-currency orgs without CurrencyIsoCode grouping.

        In multi-currency Salesforce orgs, summing Amount across records with different
        currencies produces meaningless totals (e.g., 100 USD + 50 EUR = 150 ???).

        Returns auto-fix suggestions if the issue is detected.
        """
        suggestions: List[PlanAlternative] = []

        # Check if this is a multi-currency org by looking for CurrencyIsoCode field
        primary_lower = primary_object.lower()
        schema_obj = (schema_context.objects_by_name or {}).get(primary_lower)
        if not schema_obj:
            return suggestions

        has_currency_field = "currencyisocode" in {
            f.lower() for f in schema_obj.fields.keys()
        }
        if not has_currency_field:
            return suggestions  # Not a multi-currency org

        # Check for ConvertedAmount field availability
        has_converted_amount = "convertedamount" in {
            f.lower() for f in schema_obj.fields.keys()
        }

        # Detect SUM/AVG on Amount-like fields
        amount_agg_pattern = re.compile(
            r"(?i)^(?:sum|avg)\s*\(\s*([^)]+\.)?(?:amount|arr|acv|mrr|tcv|bookings)\s*\)",
        )
        amount_aggregations: List[str] = []
        for field in plan.fields or []:
            if amount_agg_pattern.match(field):
                amount_aggregations.append(field)

        if not amount_aggregations:
            return suggestions  # No amount aggregations

        # Check if CurrencyIsoCode is already in GROUP BY or selected fields
        currency_field = f"{primary_object}.CurrencyIsoCode"
        currency_in_groupby = any(
            g.lower() == currency_field.lower() or g.lower() == "currencyisocode"
            for g in (plan.group_by or [])
        )
        currency_in_fields = any(
            f.lower() == currency_field.lower() or f.lower() == "currencyisocode"
            for f in (plan.fields or [])
        )

        if currency_in_groupby or currency_in_fields:
            return suggestions  # Already handling currency properly

        # Generate auto-fix suggestions
        agg_list = ", ".join(amount_aggregations[:2])
        if len(amount_aggregations) > 2:
            agg_list += f" (+{len(amount_aggregations) - 2} more)"

        # Suggestion 1: Group by CurrencyIsoCode
        group_ops = [
            PlanPatchOp(op="add", path="/fields/-", value=currency_field),
        ]
        if plan.group_by:
            group_ops.append(
                PlanPatchOp(op="add", path="/groupBy/-", value=currency_field)
            )
        suggestions.append(
            PlanAlternative(
                title="Group by Currency",
                description=f"Add CurrencyIsoCode to see {agg_list} broken out by currency",
                patch=_ensure_patch_summary(PlanPatch(ops=group_ops)),
            )
        )

        # Suggestion 2: Use ConvertedAmount (if available)
        if has_converted_amount:
            convert_ops: List[PlanPatchOp] = []
            for idx, field in enumerate(plan.fields or []):
                if amount_agg_pattern.match(field):
                    # Replace Amount with ConvertedAmount
                    converted = re.sub(
                        r"(?i)(\.)(amount)(\s*\))",
                        r"\1ConvertedAmount\3",
                        field,
                    )
                    if converted != field:
                        convert_ops.append(
                            PlanPatchOp(
                                op="replace", path=f"/fields/{idx}", value=converted
                            )
                        )
            if convert_ops:
                suggestions.append(
                    PlanAlternative(
                        title="Use Converted Amount",
                        description="Switch to ConvertedAmount (normalized to org default currency)",
                        patch=_ensure_patch_summary(PlanPatch(ops=convert_ops)),
                    )
                )

        return suggestions

    def _enforce_aggregation_rules(self, plan: SalesforcePlan) -> SalesforcePlan:
        """Ensure aggregate plans are SOQL-valid by aligning fields with GROUP BY."""
        agg_pattern = re.compile(
            r"(?i)^(?P<func>sum|avg|count|min|max)\s*\((?P<inner>.+)\)\s*$"
        )
        bool_suffix_pattern = re.compile(r"(?i)(?:^|\.)(iswon|isclosed)$")
        metric_keywords = (
            "amount",
            "revenue",
            "arr",
            "mrr",
            "tcv",
            "acv",
            "bookings",
            "probability",
            "expected",
            "quota",
            "target",
            "score",
            "number",
            "convertedamount",
        )
        dimension_hints = (
            "name",
            "stage",
            "status",
            "type",
            "category",
            "source",
            "owner",
            "industry",
            "segment",
            "region",
            "territory",
            "fiscal",
            "quarter",
            "month",
            "year",
            "currencyisocode",
        )

        def is_metric_field(field_ref: str) -> bool:
            last = (field_ref or "").split(".")[-1].lower()
            if last.endswith("id"):
                return True
            if any(hint in last for hint in dimension_hints):
                return False
            return any(keyword in last for keyword in metric_keywords)

        parsed_fields: List[Tuple[str, str, str, bool]] = []
        has_aggregate = False
        for field in plan.fields or []:
            raw_field = field.strip()
            base_field = raw_field
            alias = None
            alias_parts = re.split(r"(?i)\s+as\s+", raw_field, maxsplit=1)
            if len(alias_parts) == 2:
                base_field = alias_parts[0].strip()
                alias = alias_parts[1].strip()
            agg_match = agg_pattern.match(base_field)
            if agg_match:
                has_aggregate = True
                func = agg_match.group("func").upper()
                inner = agg_match.group("inner").strip()
                fixed = raw_field
                if func in {"AVG", "SUM"} and bool_suffix_pattern.search(inner):
                    replacement = re.sub(
                        r"(?i)(iswon|isclosed)$", "Probability", inner
                    )
                    fixed = f"{func}({replacement})"
                    if alias:
                        fixed = f"{fixed} AS {alias}"
                    if fixed != field:
                        log.info(
                            "Replacing boolean aggregate '%s' with '%s'",
                            field,
                            fixed,
                        )
                parsed_fields.append((raw_field, base_field, fixed, True))
            else:
                parsed_fields.append((raw_field, base_field, raw_field, False))

        if not has_aggregate and not plan.group_by:
            return plan

        if has_aggregate and not plan.group_by:
            group_candidates: List[str] = []
            for raw_field, base_field, _, is_agg in parsed_fields:
                if is_agg:
                    continue
                if is_metric_field(base_field):
                    continue
                group_candidates.append(base_field)
            if group_candidates:
                plan.group_by = list(dict.fromkeys(group_candidates))
                log.info(
                    "Added GROUP BY fields to satisfy aggregation rules: %s",
                    plan.group_by,
                )

        safe_fields = {field.lower() for field in plan.group_by or []}
        cleaned_fields: List[str] = []
        for raw_field, base_field, fixed, is_agg in parsed_fields:
            if is_agg:
                cleaned_fields.append(fixed)
                continue
            if base_field.lower() in safe_fields:
                cleaned_fields.append(raw_field)
                continue
            if has_aggregate or plan.group_by:
                log.info(
                    "Dropping invalid raw field '%s' from aggregated plan",
                    raw_field,
                )
        plan.fields = list(dict.fromkeys(cleaned_fields))
        return plan

    def _resolve_primary_object(self, plan: SalesforcePlan) -> str:
        """Promote child objects to primary for aggregate queries.

        In SOQL, you must query FROM the child object when aggregating child fields.
        E.g., "Which campaigns generated the most pipeline?" should query:
            FROM Opportunity ... GROUP BY Campaign.Name
        Not:
            FROM Campaign WHERE Opportunity.CloseDate = ...
        """
        if not plan.objects:
            return "Lead"

        current_primary = plan.objects[0]

        # Only apply promotion when we have aggregates
        has_aggregates = any("(" in f for f in (plan.fields or []))
        if not has_aggregates:
            return current_primary

        # Child objects that typically hold aggregatable data (Amount, counts, etc.)
        heavy_child_objects = {"opportunity", "case", "opportunitylineitem", "task", "event"}

        # Check if any heavy child is in the objects list but not primary
        for obj in plan.objects:
            if obj.lower() in heavy_child_objects and obj.lower() != current_primary.lower():
                log.info(
                    "Promoting %s to Primary Object for semantic correctness (was %s)",
                    obj,
                    current_primary,
                )
                return obj

        return current_primary

    def compile_plan(self, plan: SalesforcePlan) -> str:
        self._translate_filters(plan)
        for obj_name in plan.objects or []:
            _validate_soql_identifier(obj_name, "object")
        primary = self._resolve_primary_object(plan)
        _validate_soql_identifier(primary, "object")
        select_fields = plan.fields or [f"{primary}.Id"]
        normalized_fields: List[str] = []
        for field in select_fields:
            normalized = _validate_soql_field_expr(field, "field")
            if _SOQL_AGGREGATE_FIELD_PATTERN.fullmatch(normalized):
                normalized_fields.append(normalized)  # Keep aggregates as-is
            elif "." not in normalized:
                normalized_fields.append(f"{primary}.{normalized}")
            else:
                normalized_fields.append(normalized)

        select_clause = ", ".join(normalized_fields)
        soql = [f"SELECT {select_clause}", f"FROM {primary}"]

        if plan.filters:
            where_clauses: List[str] = []
            for filt in plan.filters:
                field_expr = _validate_soql_field_expr(filt.field, "filter field")
                op = filt.op
                if op in {"IN", "NOT_IN"}:
                    values = filt.values or []
                    formatted = ", ".join(_format_value(v) for v in values[:25])
                    operator = "IN" if op == "IN" else "NOT IN"
                    where_clauses.append(f"{field_expr} {operator} ({formatted})")
                elif op == "DATE_LITERAL":
                    date_value = str(filt.value or "").strip()
                    is_named_literal = bool(
                        _SOQL_DATE_LITERAL_PATTERN.fullmatch(date_value)
                    )
                    is_iso_date = bool(
                        _SOQL_STRICT_DATE_VALUE_PATTERN.fullmatch(date_value)
                    )
                    if not (is_named_literal or is_iso_date):
                        raise PlanValidationError(
                            f"Invalid DATE literal '{date_value}'"
                        )
                    where_clauses.append(f"{field_expr} = {date_value}")
                elif op == "IS_NULL":
                    where_clauses.append(f"{field_expr} = null")
                elif op == "IS_NOT_NULL":
                    where_clauses.append(f"{field_expr} != null")
                elif op == "BETWEEN":
                    values = filt.values or []
                    if len(values) >= 2:
                        where_clauses.append(
                            f"{field_expr} >= {_format_value(values[0])} AND {field_expr} <= {_format_value(values[1])}"
                        )
                elif op in {"CONTAINS", "STARTS_WITH", "ENDS_WITH", "LIKE"}:
                    raw_value = str(filt.value or "")
                    escaped_value = _escape_soql_literal(raw_value)
                    if op == "CONTAINS":
                        pattern = f"'%{escaped_value}%'"
                    elif op == "STARTS_WITH":
                        pattern = f"'{escaped_value}%'"
                    elif op == "ENDS_WITH":
                        pattern = f"'%{escaped_value}'"
                    else:
                        pattern = _format_value(raw_value)
                    where_clauses.append(f"{field_expr} LIKE {pattern}")
                else:
                    where_clauses.append(
                        f"{field_expr} {op} {_format_value(filt.value)}"
                    )
            if where_clauses:
                soql.append("WHERE " + " AND ".join(where_clauses))

        if plan.group_by:
            validated_group_by = [
                _validate_soql_field_expr(group_field, "group by field")
                for group_field in plan.group_by
            ]
            soql.append("GROUP BY " + ", ".join(validated_group_by))

        if plan.order_by:
            order_clause = ", ".join(
                f"{_validate_soql_field_expr(item.field, 'order by field')} {item.dir}"
                for item in plan.order_by
            )
            soql.append("ORDER BY " + order_clause)

        if plan.limit:
            soql.append(f"LIMIT {plan.limit}")

        return " ".join(soql)

    async def preview_plan(
        self,
        *,
        plan: SalesforcePlan,
        gateway: SalesforceGateway,
        account_id: str,
        tier: Optional[str],
        schema_context: Optional[SchemaPromptContext] = None,
        request_id: Optional[str] = None,
    ) -> Dict[str, Any]:
        if schema_context is None:
            raise PlanValidationError(
                "Schema context is required for preview allowlist enforcement"
            )
        preview_plan = plan.model_copy(deep=True)
        preview_plan.limit = min(
            AI_REPORTS_PREVIEW_LIMIT, preview_plan.limit or AI_REPORTS_PREVIEW_LIMIT
        )
        soql = self.compile_plan(preview_plan)
        allowlist = _collect_allowlist(preview_plan, schema_context=schema_context)
        start = time.perf_counter()
        result = await gateway.soql(
            soql,
            admin_override=False,
            user=f"ai-reports@{account_id}",
            ip="0.0.0.0",
            user_is_admin=False,
            plan_allowlist=allowlist,
        )
        duration_ms = int((time.perf_counter() - start) * 1000)
        records = result.get("records", []) or []
        total_size = result.get("totalSize", len(records))
        try:
            emit(
                "preview.outcome",
                request_id=request_id,
                org_id=account_id,
                ms=duration_ms,
                rows=len(records),
                total=total_size,
            )
        except Exception:  # noqa: BLE001
            log.debug("Failed to emit preview telemetry", exc_info=True)
        if not records:
            return {"headers": [], "rows": [], "soql": soql, "totalSize": total_size}

        headers = [key for key in records[0].keys() if key != "attributes"]

        def normalize_cell(cell: Any) -> Any:
            if cell is None:
                return ""
            if isinstance(cell, dict):
                if "Name" in cell:
                    return cell["Name"]
                if "name" in cell:
                    return cell["name"]
                return json.dumps(cell, ensure_ascii=False)
            return cell

        rows = [
            [normalize_cell(row.get(header)) for header in headers] for row in records
        ]
        return {
            "headers": headers,
            "rows": rows,
            "soql": soql,
            "totalSize": total_size,
        }

    async def _repair_plan(
        self,
        *,
        plan: SalesforcePlan,
        error_message: str,
        error_type: str,
        invalid_fields: List[str],
        schema_context: SchemaPromptContext,
        tier: Optional[str],
        request_id: Optional[str] = None,
    ) -> Tuple[SalesforcePlan, List[str]]:
        """
        Re-run planner with error context to fix invalid fields/relationships.

        Uses HIGH thinking level for deep reasoning about the error.
        Returns (repaired_plan, repair_trace).
        """
        # Build a minimal schema subset focusing on the problem areas
        # Extract object names from invalid fields for targeted context
        problem_objects: Set[str] = set()
        for field in invalid_fields:
            if "." in field:
                problem_objects.add(field.split(".")[0])
        # Also include all objects from the original plan
        for obj in plan.objects:
            problem_objects.add(obj)

        # Build repair prompt with error context
        repair_instructions = [
            f"The previous query plan failed with error type: {error_type}",
            f"Error message: {error_message}",
        ]

        if invalid_fields:
            repair_instructions.append(
                f"Invalid fields/relationships that must be removed or corrected: {', '.join(invalid_fields)}"
            )

        if error_type == "FIXABLE_AGGREGATION":
            repair_instructions.append(
                "AGGREGATION FIX: Ensure all non-aggregated fields appear in GROUP BY, "
                "or wrap them in aggregate functions (COUNT, SUM, AVG, MAX, MIN)."
            )
        elif error_type == "FIXABLE_RELATIONSHIP":
            repair_instructions.append(
                "RELATIONSHIP FIX: Use only valid parent-child relationships. "
                "Check that relationship names match the Salesforce schema exactly."
            )
        elif error_type == "FIXABLE_SYNTAX":
            repair_instructions.append(
                "SYNTAX FIX: Review SOQL syntax rules. Ensure proper quoting, "
                "valid operators, and correct date literal formats."
            )

        repair_instructions.append(
            "Generate a corrected plan using ONLY valid fields from the schema. "
            "If a requested field does not exist, omit it or find a valid alternative."
        )

        repair_prompt = "\n".join(repair_instructions)

        # Log repair attempt
        log.info(
            "Attempting plan repair: type=%s, invalid_fields=%s, request_id=%s",
            error_type,
            invalid_fields,
            request_id,
        )

        try:
            draft = await self._run_planner(
                prompt=repair_prompt,
                schema_context=schema_context,
                research=None,  # Skip research phase for repair
                tier=tier,
                prior_warnings=[error_message],
            )

            repair_trace = [
                f"Repair attempt ({error_type}): addressed {', '.join(invalid_fields) if invalid_fields else 'syntax/structure issues'}"
            ]

            # Emit telemetry for repair attempt
            try:
                emit(
                    "repair.attempt",
                    request_id=request_id,
                    error_type=error_type,
                    invalid_fields=invalid_fields,
                    success=True,
                )
            except Exception:
                pass

            return draft.plan, repair_trace

        except Exception as repair_exc:
            log.warning(
                "Plan repair failed: %s (original error: %s)",
                repair_exc,
                error_message,
            )
            try:
                emit(
                    "repair.attempt",
                    request_id=request_id,
                    error_type=error_type,
                    invalid_fields=invalid_fields,
                    success=False,
                    repair_error=str(repair_exc),
                )
            except Exception:
                pass
            raise

    async def preview_plan_with_repair(
        self,
        *,
        plan: SalesforcePlan,
        gateway: SalesforceGateway,
        account_id: str,
        tier: Optional[str],
        schema_context: SchemaPromptContext,
        request_id: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Execute preview with automatic repair on repairable SOQL errors.

        Attempts up to MAX_REPAIR_ATTEMPTS repairs for fixable errors
        (invalid fields, relationships, aggregation, syntax).
        Non-fixable errors (permissions, timeouts) are raised immediately.
        """
        current_plan = plan
        repair_trace: List[str] = []

        for attempt in range(MAX_REPAIR_ATTEMPTS + 1):
            try:
                result = await self.preview_plan(
                    plan=current_plan,
                    gateway=gateway,
                    account_id=account_id,
                    tier=tier,
                    schema_context=schema_context,
                    request_id=request_id,
                )
                # Add repair trace to result if any repairs were made
                if repair_trace:
                    result["repair_trace"] = repair_trace
                return result

            except Exception as exc:
                error_message = str(exc)

                # Check if this error is repairable
                if not _is_repairable_error(error_message):
                    log.debug(
                        "Non-repairable error, raising immediately: %s",
                        error_message[:200],
                    )
                    raise

                # Check if we've exhausted repair attempts
                if attempt >= MAX_REPAIR_ATTEMPTS:
                    log.warning(
                        "Max repair attempts (%d) exceeded for request_id=%s",
                        MAX_REPAIR_ATTEMPTS,
                        request_id,
                    )
                    try:
                        emit(
                            "repair.exhausted",
                            request_id=request_id,
                            attempts=attempt,
                            final_error=error_message[:500],
                        )
                    except Exception:
                        pass
                    raise

                # Classify and extract error details
                error_type = _classify_soql_error(error_message)
                invalid_fields = _extract_invalid_fields(error_message)

                log.info(
                    "Repairable error detected (attempt %d/%d): type=%s, fields=%s",
                    attempt + 1,
                    MAX_REPAIR_ATTEMPTS,
                    error_type,
                    invalid_fields,
                )

                # Attempt repair
                try:
                    current_plan, trace = await self._repair_plan(
                        plan=current_plan,
                        error_message=error_message,
                        error_type=error_type,
                        invalid_fields=invalid_fields,
                        schema_context=schema_context,
                        tier=tier,
                        request_id=request_id,
                    )
                    repair_trace.extend(trace)
                except Exception as repair_exc:
                    # If repair itself fails, raise the original error
                    log.warning(
                        "Repair attempt failed, raising original error: %s",
                        repair_exc,
                    )
                    raise exc from repair_exc

        # Should not reach here, but raise if we do
        raise RuntimeError("Unexpected exit from repair loop")


def hash_plan_for_storage(plan: SalesforcePlan) -> str:
    return _hash_plan_dict(plan.model_dump(by_alias=True))
