"""Architect Service - Gemini Flash converts intent to FlowPlan JSON."""

from __future__ import annotations

import logging
import os
import re
from dataclasses import dataclass
from typing import Any, Dict, List, Optional

from pydantic import ValidationError

from ..api.v2.flow_builder_schemas import (
    ConditionOperator,
    FlowAudit,
    FlowAuditPatch,
    FlowAuditRisk,
    FlowAuditSuggestion,
    FlowCondition,
    FlowEntryCriteria,
    FlowPlan,
    FlowTriggerType,
    RiskScore,
    SalesforceFieldType,
    ScoutContext,
)
from .. import settings
from .gemini_client import GeminiClient
from .flow_async_fields import collect_condition_fields, get_async_field_category
from .llm_guard import ensure_routing_snapshot, save_routing_snapshot

log = logging.getLogger(__name__)

_FALLBACK_ERROR_TYPES = {
    "ResourceExhausted",
    "ServiceUnavailable",
    "DeadlineExceeded",
    "TimeoutError",
    "LLMCircuitOpenError",
}
_FALLBACK_PATTERNS = (
    "429",
    "rate limit",
    "ratelimit",
    "quota exceeded",
    "resource exhausted",
    "service unavailable",
    "503",
    "timeout",
    "deadline exceeded",
)


def _should_fallback(raw: Dict[str, Any]) -> bool:
    if not raw or not raw.get("_parse_error"):
        return False
    err_type = str(raw.get("_error_type") or "")
    if err_type in _FALLBACK_ERROR_TYPES:
        return True
    if err_type == "LLMConcurrencyError":
        return False
    haystack = f"{raw.get('_error') or ''} {raw.get('_raw') or ''}".lower()
    return any(token in haystack for token in _FALLBACK_PATTERNS)


ARCHITECT_SYSTEM_PROMPT = """You are the Architect for FoundryOps Flow Builder.

Convert user intent into FlowPlan JSON. Output ONLY valid JSON.

## Output Schema
{
  "flow_name": "FO_<PascalCaseName>",
  "description": "Plain English description",
  "trigger": {
    "object_api_name": "Lead",
    "type": "RecordBeforeSave",
    "events": ["Create"] | ["Update"] | ["Create", "Update"],
    "record_type_id": null | "012..."
  },
  "entry_criteria": {
    "logic": "AND" | "OR",
    "conditions": [{
      "field": "FieldApiName",
      "operator": "EqualTo|NotEqualTo|IsNull|IsNotNull|Contains|StartsWith|GreaterThan|LessThan",
      "value": "string|number|boolean|null",
      "value_is_field": false
    }]
  },
  "branches": [{
    "name": "Working when email sent",
    "logic": "AND" | "OR",
    "conditions": [{
      "field": "FieldApiName",
      "operator": "EqualTo|NotEqualTo|IsNull|IsNotNull|Contains|StartsWith|GreaterThan|LessThan",
      "value": "string|number|boolean|null",
      "value_is_field": false
    }],
    "actions": [{
      "type": "assignment",
      "field": "TargetField",
      "value": {"source": "literal|field_reference|formula", "content": "value"}
    }]
  }],
  "actions": [{
    "type": "assignment",
    "field": "TargetField",
    "value": {"source": "literal|field_reference|formula", "content": "value"}
  }],
  "safety_check": {"risk_score": "Low|Medium|High", "reasoning": "explanation"}
}

## V1 Constraints (MUST FOLLOW)
- trigger.type MUST be "RecordBeforeSave" unless the conditions reference EAC_* fields, in which case use "RecordAfterSave"
- trigger.events MUST be ["Create"], ["Update"], or ["Create", "Update"]
- logic: only "AND" or "OR" (no CUSTOM)
- actions: only "assignment" type
- Max 10 conditions, 20 actions
- Max 5 branches, one level only (no nesting)
- flow_name MUST start with "FO_"

## CRITICAL SAFETY RULES
1. If events include "Update", you MUST include either entry_criteria with at least one condition OR branch conditions
2. If using branches, you MUST NOT include both actions and branches
3. If using branches, each branch MUST have its own conditions
4. Default to "set if blank": For non-boolean fields, add IsNull condition unless user explicitly wants overwrite
5. Use exact field API names from context (case-sensitive)
6. For picklists, use exact values from metadata
7. EAC_* activity fields are not supported in RecordBeforeSave flows; use RecordAfterSave when they are used in conditions

## Default Behavior (unless user explicitly requests otherwise)
 - When setting a field, add a condition (entry_criteria or branch) : TargetField IsNull
 - This prevents overwriting existing values

## Branching Guidance
- Use "branches" only when the user wants different actions under different conditions.
- No nested decisions. No default/else branch. If nothing matches, do nothing.
- If branching is requested, set actions=[] (or omit) and populate branches instead.

Return ONLY JSON, no explanation."""

CRITIC_SYSTEM_PROMPT = """You are the Critic auditing a Salesforce Flow plan.

You provide ADVISORY feedback. Sheriff is the actual deployment blocker.

## Output Schema (FlowAudit)
{
  "risk_score": "Low|Medium|High|Critical",
  "risks": [{
    "severity": "warning|error",
    "category": "data_loss|infinite_loop|validation_conflict|field_access|logic_error",
    "message": "description",
    "field": "field_name or null"
  }],
  "suggestions": [{
    "title": "short title",
    "description": "what to fix"
  }],
  "patches": [{
    "title": "short title",
    "description": "what to change",
    "patch": "optional patch text"
  }],
  "requires_acknowledgment": true|false
}

## Risk Scoring
- Low: Safe operations, good entry criteria
- Medium: Some risk but mitigated by conditions
- High: Potential data issues, needs review
- Critical: Very risky - set requires_acknowledgment=true

## Critical Patterns to Flag
1. INFINITE LOOP: Updates field in entry criteria without guard
2. DATA LOSS: Overwrites non-null fields unconditionally
3. VALIDATION CONFLICT: Sets value that may violate validation rules
4. FIELD ACCESS: References non-updateable or formula field

NOTE: You cannot block deployment. Set requires_acknowledgment=true for Critical risks."""


@dataclass
class FlowArchitectResult:
    plan: Optional[FlowPlan]
    audit: Optional[FlowAudit]
    clarifiers: List[str]
    raw: Dict[str, Any]
    model_used: str
    error: Optional[str] = None


class FlowArchitect:
    """Gemini Flash-powered architect that outputs FlowPlan JSON."""

    def __init__(
        self,
        model: Optional[str] = None,
        timeout_s: Optional[int] = None,
    ) -> None:
        resolved_model = (
            model
            or os.getenv("FLOW_ARCHITECT_MODEL")
            or os.getenv("GEMINI_MODEL")
            or "gemini-3-flash-preview"
        )
        max_tokens = int(os.getenv("FLOW_ARCHITECT_MAX_OUTPUT_TOKENS", "8192"))
        timeout_s = timeout_s or settings.LLM_HEAVY_REASON_TIMEOUT_S
        self._timeout_s = timeout_s
        self._max_tokens = max_tokens
        self._max_retries = 2
        self._budget_s = settings.LLM_HEAVY_REASON_BUDGET_S
        self._retry_jitter_min_ms = settings.LLM_RETRY_JITTER_MIN_MS
        self._retry_jitter_max_ms = settings.LLM_RETRY_JITTER_MAX_MS
        self._fallback_model = settings.LLM_FALLBACK_MODEL
        self.client = GeminiClient(
            model=resolved_model,
            timeout_s=timeout_s,
            max_output_tokens=max_tokens,
            max_retries=self._max_retries,
            wall_clock_budget_s=self._budget_s,
            retry_jitter_min_ms=self._retry_jitter_min_ms,
            retry_jitter_max_ms=self._retry_jitter_max_ms,
        )
        self.model_used = resolved_model

    def _make_client(self, model: str, *, max_output_tokens: Optional[int] = None) -> GeminiClient:
        return GeminiClient(
            model=model,
            timeout_s=self._timeout_s,
            max_output_tokens=max_output_tokens or self._max_tokens,
            max_retries=self._max_retries,
            wall_clock_budget_s=self._budget_s,
            retry_jitter_min_ms=self._retry_jitter_min_ms,
            retry_jitter_max_ms=self._retry_jitter_max_ms,
        )

    def _resolve_snapshot(
        self, job_id: Optional[str], primary_model: str
    ):
        if not job_id:
            return None
        key = f"llm:job_routing:{job_id}"
        return ensure_routing_snapshot(
            key,
            primary_model=primary_model,
            fallback_model=self._fallback_model,
            ttl_s=3600,
            tier="flow_architect",
        )

    def _apply_fallback_snapshot(
        self, job_id: Optional[str], snapshot
    ) -> None:
        if not job_id or not snapshot:
            return
        key = f"llm:job_routing:{job_id}"
        save_routing_snapshot(key, snapshot, ttl_s=3600)

    def build_plan(
        self,
        *,
        prompt: str,
        object_api_name: str,
        context: Optional[ScoutContext] = None,
        record_type_id: Optional[str] = None,
        run_critic: bool = True,
        job_id: Optional[str] = None,
        account_id: Optional[str] = None,
    ) -> FlowArchitectResult:
        clarifiers = self._detect_clarifiers(prompt, context)
        user_prompt = self._build_user_prompt(
            prompt=prompt,
            object_api_name=object_api_name,
            context=context,
            record_type_id=record_type_id,
        )
        messages = [
            {"role": "system", "content": ARCHITECT_SYSTEM_PROMPT},
            {"role": "user", "content": user_prompt},
        ]

        snapshot = self._resolve_snapshot(job_id, self.model_used)
        use_fallback = bool(snapshot and snapshot.using_fallback and self._fallback_model)
        model_used = self._fallback_model if use_fallback else self.model_used
        client = self.client if not use_fallback else self._make_client(self._fallback_model)

        raw = client.chat_json(messages, retry_strict=True, account_id=account_id)
        if raw.get("_parse_error") and not raw.get("_error_type"):
            # One extra retry for robustness
            raw = client.chat_json(messages, retry_strict=True, account_id=account_id)

        if (
            raw.get("_parse_error")
            and not use_fallback
            and self._fallback_model
            and _should_fallback(raw)
        ):
            if snapshot:
                snapshot.using_fallback = True
                self._apply_fallback_snapshot(job_id, snapshot)
            fallback_client = self._make_client(self._fallback_model)
            raw = fallback_client.chat_json(
                messages, retry_strict=True, account_id=account_id
            )
            if raw.get("_parse_error") and not raw.get("_error_type"):
                raw = fallback_client.chat_json(
                    messages, retry_strict=True, account_id=account_id
                )
            model_used = self._fallback_model

        if raw.get("_parse_error"):
            return FlowArchitectResult(
                plan=None,
                audit=None,
                clarifiers=clarifiers,
                raw=raw,
                model_used=model_used,
                error=raw.get("_error") or "json_parse_error",
            )

        try:
            plan = FlowPlan.model_validate(raw)
        except ValidationError as exc:
            return FlowArchitectResult(
                plan=None,
                audit=None,
                clarifiers=clarifiers,
                raw=raw,
                model_used=model_used,
                error=str(exc),
            )

        plan = self._apply_safe_defaults(plan, context, prompt)
        plan = self._ensure_supported_trigger(plan)
        if record_type_id and not plan.trigger.record_type_id:
            plan = plan.model_copy(
                update={
                    "trigger": plan.trigger.model_copy(
                        update={"record_type_id": record_type_id}
                    )
                }
            )
        if not plan.user_prompt:
            plan = plan.model_copy(update={"user_prompt": prompt})

        audit = (
            self.audit_plan(plan, context, job_id=job_id, account_id=account_id)
            if run_critic
            else None
        )

        return FlowArchitectResult(
            plan=plan,
            audit=audit,
            clarifiers=clarifiers,
            raw=raw,
            model_used=model_used,
        )

    def audit_plan(
        self,
        plan: FlowPlan,
        context: Optional[ScoutContext] = None,
        *,
        job_id: Optional[str] = None,
        account_id: Optional[str] = None,
    ) -> FlowAudit:
        critic_model = os.getenv("FLOW_CRITIC_MODEL") or "gemini-3-pro-preview"
        critic_max_tokens = int(os.getenv("FLOW_CRITIC_MAX_OUTPUT_TOKENS", "2048"))
        snapshot = self._resolve_snapshot(job_id, critic_model)
        use_fallback = bool(snapshot and snapshot.using_fallback and self._fallback_model)
        resolved_model = self._fallback_model if use_fallback else critic_model
        critic = self._make_client(resolved_model, max_output_tokens=critic_max_tokens)

        prompt = self._build_critic_prompt(plan, context)
        raw = critic.chat_json(
            [
                {"role": "system", "content": CRITIC_SYSTEM_PROMPT},
                {"role": "user", "content": prompt},
            ],
            retry_strict=True,
            account_id=account_id,
        )
        if raw.get("_parse_error") and not raw.get("_error_type"):
            raw = critic.chat_json(
                [
                    {"role": "system", "content": CRITIC_SYSTEM_PROMPT},
                    {"role": "user", "content": prompt},
                ],
                retry_strict=True,
                account_id=account_id,
            )

        if (
            raw.get("_parse_error")
            and not use_fallback
            and self._fallback_model
            and _should_fallback(raw)
        ):
            if snapshot:
                snapshot.using_fallback = True
                self._apply_fallback_snapshot(job_id, snapshot)
            fallback_client = self._make_client(
                self._fallback_model, max_output_tokens=critic_max_tokens
            )
            raw = fallback_client.chat_json(
                [
                    {"role": "system", "content": CRITIC_SYSTEM_PROMPT},
                    {"role": "user", "content": prompt},
                ],
                retry_strict=True,
                account_id=account_id,
            )
            if raw.get("_parse_error") and not raw.get("_error_type"):
                raw = fallback_client.chat_json(
                    [
                        {"role": "system", "content": CRITIC_SYSTEM_PROMPT},
                        {"role": "user", "content": prompt},
                    ],
                    retry_strict=True,
                    account_id=account_id,
                )

        if raw.get("_parse_error"):
            return FlowAudit(
                risk_score=RiskScore.MEDIUM,
                risks=[
                    FlowAuditRisk(
                        severity="warning",
                        category="unknown",
                        message="Critic failed to parse response; review manually.",
                        field=None,
                    )
                ],
                suggestions=[],
                patches=[],
                requires_acknowledgment=False,
            )

        try:
            audit = FlowAudit.model_validate(raw)
            if audit.risk_score == RiskScore.CRITICAL and not audit.requires_acknowledgment:
                audit = audit.model_copy(update={"requires_acknowledgment": True})
            return audit
        except ValidationError:
            return FlowAudit(
                risk_score=RiskScore.MEDIUM,
                risks=[
                    FlowAuditRisk(
                        severity="warning",
                        category="unknown",
                        message="Critic response failed validation; review manually.",
                        field=None,
                    )
                ],
                suggestions=[],
                patches=[],
                requires_acknowledgment=False,
            )

    @staticmethod
    def _ensure_supported_trigger(plan: FlowPlan) -> FlowPlan:
        """Switch to RecordAfterSave when conditions reference async-populated fields."""
        if not plan or not plan.trigger:
            return plan
        if plan.trigger.type == FlowTriggerType.RECORD_AFTER_SAVE:
            return plan

        condition_fields: List[str] = []
        if plan.entry_criteria:
            condition_fields.extend(
                list(collect_condition_fields(plan.entry_criteria.conditions))
            )
        if plan.branches:
            for branch in plan.branches:
                condition_fields.extend(
                    list(collect_condition_fields(branch.conditions or []))
                )

        if any(get_async_field_category(f) for f in condition_fields):
            return plan.model_copy(
                update={
                    "trigger": plan.trigger.model_copy(
                        update={"type": FlowTriggerType.RECORD_AFTER_SAVE}
                    )
                }
            )

        return plan

    def _build_user_prompt(
        self,
        *,
        prompt: str,
        object_api_name: str,
        context: Optional[ScoutContext],
        record_type_id: Optional[str],
    ) -> str:
        sections = [
            f"Object: {object_api_name}",
            f"User Request: {prompt}",
        ]
        if record_type_id:
            sections.append(f"Selected RecordTypeId: {record_type_id}")

        if context:
            sections.append(self._format_context(context, record_type_id))
        return "\n\n".join(sections)

    def _format_context(
        self, context: ScoutContext, record_type_id: Optional[str]
    ) -> str:
        lines: List[str] = []
        lines.append(f"Object Label: {context.object_label}")

        if context.record_types:
            rt_lines = []
            for rt in context.record_types[:10]:
                flag = " (default)" if rt.is_default else ""
                rt_lines.append(f"{rt.name} [{rt.id}]{flag}")
            lines.append("Record Types: " + "; ".join(rt_lines))

        picklists = []
        for name, field in context.fields.items():
            if field.type in (
                SalesforceFieldType.PICKLIST,
                SalesforceFieldType.MULTIPICKLIST,
            ):
                values = context.get_picklist_values(name, record_type_id)
                if values:
                    picklists.append(f"- {name}: {values[:12]}")
                else:
                    picklists.append(f"- {name}: (values require record type)")
        if picklists:
            lines.append("Picklist Values:\n" + "\n".join(picklists[:50]))

        field_lines = []
        for idx, field in enumerate(context.fields.values()):
            if idx >= 200:
                field_lines.append("... (truncated)")
                break
            field_lines.append(f"- {field.name} ({field.type})")
        lines.append("Fields:\n" + "\n".join(field_lines))

        return "\n".join(lines)

    def _build_critic_prompt(
        self, plan: FlowPlan, context: Optional[ScoutContext]
    ) -> str:
        import json

        plan_json = json.dumps(plan.model_dump(mode="json"), sort_keys=True)
        sections = [
            "FlowPlan JSON:",
            f"{plan_json}",
        ]
        if context:
            rules = []
            for rule in context.validation_rules[:20]:
                rules.append(
                    f"- {rule.name}: {rule.error_condition_formula or ''} ({rule.error_message or ''})"
                )
            if rules:
                sections.append("Validation Rules (truncated):\n" + "\n".join(rules))
        return "\n\n".join(sections)

    def _detect_clarifiers(
        self, prompt: str, context: Optional[ScoutContext]
    ) -> List[str]:
        clarifiers: List[str] = []
        low = (prompt or "").lower()

        if context and context.fields:
            field_mentions = [
                name
                for name in context.fields.keys()
                if name.lower() in low
            ]
            if not field_mentions:
                clarifiers.append("Which field(s) should be updated?")

        if not any(
            token in low
            for token in ("create", "update", "when ", "on ", "before save")
        ):
            clarifiers.append(
                "Should this run on record Create, Update, or both?"
            )

        if any(
            token in low
            for token in ("set", "update", "assign", "copy", "fill", "sync")
        ):
            if not re.search(r"\b(to|from|=)\b", low):
                clarifiers.append("What value should be set, or copied from which field?")

        # Cap clarifiers to keep UI simple
        return clarifiers[:3]

    def _apply_safe_defaults(
        self, plan: FlowPlan, context: Optional[ScoutContext], prompt: str
    ) -> FlowPlan:
        if not context or self._wants_overwrite(prompt):
            return plan

        if plan.branches:
            updated_branches = []
            for branch in plan.branches:
                if branch.logic == "OR":
                    updated_branches.append(branch)
                    continue

                existing_fields = {cond.field for cond in branch.conditions}
                conditions: List[FlowCondition] = []
                for action in branch.actions:
                    field_meta = context.get_field(action.field)
                    if not field_meta:
                        continue
                    if field_meta.type == SalesforceFieldType.BOOLEAN:
                        continue
                    if action.field in existing_fields:
                        continue
                    if action.value.content in (None, ""):
                        continue
                    conditions.append(
                        FlowCondition(
                            field=action.field,
                            operator=ConditionOperator.IS_NULL,
                            value=None,
                        )
                    )

                if conditions:
                    branch = branch.model_copy(
                        update={"conditions": branch.conditions + conditions}
                    )
                updated_branches.append(branch)

            return plan.model_copy(update={"branches": updated_branches})

        entry = plan.entry_criteria
        if entry and entry.logic == "OR":
            return plan

        existing_fields = set()
        if entry:
            for cond in entry.conditions:
                existing_fields.add(cond.field)

        conditions = []
        for action in plan.actions:
            field_meta = context.get_field(action.field)
            if not field_meta:
                continue
            if field_meta.type == SalesforceFieldType.BOOLEAN:
                continue
            if action.field in existing_fields:
                continue
            if action.value.content in (None, ""):
                continue
            conditions.append(
                FlowCondition(
                    field=action.field,
                    operator=ConditionOperator.IS_NULL,
                    value=None,
                )
            )

        if not conditions:
            return plan

        if entry:
            entry = entry.model_copy(update={"conditions": entry.conditions + conditions})
        else:
            entry = FlowEntryCriteria(logic="AND", conditions=conditions)

        return plan.model_copy(update={"entry_criteria": entry})

    @staticmethod
    def _wants_overwrite(prompt: str) -> bool:
        low = (prompt or "").lower()
        overwrite_markers = (
            "overwrite",
            "over-write",
            "replace",
            "always set",
            "always update",
            "force",
            "regardless",
            "even if",
            "no matter what",
            "set every time",
            "update every time",
            "clear the field",
            "blank out",
            "reset the field",
        )
        return any(marker in low for marker in overwrite_markers)
