# domain/proposal_engine.py - V7.2 (Deterministic Governed Reasoning Runtime)
import json
import logging
from typing import List
from enum import Enum
from pydantic import BaseModel, ValidationError
from .ontology import get_allowed_concepts
import domain.telemetry as telemetry

logger = logging.getLogger(__name__)

# ==================== V7.2: STRICT SCHEMA DEFINITIONS ====================

class ComputeAction(str, Enum):
    """V7.3: Hardcoded allowed actions. The Planner may NOT invent new ones."""
    # ── Algebraic (V7.2) ──────────────────────────────────────────────────
    SOLVE_EQUATION    = "SOLVE_EQUATION"
    SIMPLIFY          = "SIMPLIFY"
    FIND_DERIVATIVE   = "FIND_DERIVATIVE"
    FIND_INTEGRAL     = "FIND_INTEGRAL"
    FACTOR            = "FACTOR"
    EXPAND            = "EXPAND"
    SUBSTITUTE        = "SUBSTITUTE"
    # ── Geometry / Analytic (V7.3) ──────────────────────────────────────
    FIND_AXIS_INTERSECTIONS  = "FIND_AXIS_INTERSECTIONS"   # הצב x=0 / y=0 ופתור
    CALCULATE_SLOPE_AND_LINE = "CALCULATE_SLOPE_AND_LINE"  # שיפוע + משואת ישר
    CALCULATE_DISTANCE       = "CALCULATE_DISTANCE"        # נוסחת מרחק + בדיקת מיקום

class ServerComputeTask(BaseModel):
    action: ComputeAction         # Enum enforced: unknown action → ValidationError
    target_step_ref: str          # AST Node ID only (e.g. "ast_node_2"), NOT a math expression

class PlannerResponse(BaseModel):
    pedagogical_rationale: str
    requires_server_compute: List[ServerComputeTask]

class PlannerFormatError(Exception):
    """Raised when the Planner's response fails structural or schema validation."""
    pass

MAX_SERVER_STEPS = 8  # DOS Prevention: Planner cannot request more than 8 server tasks


# ==================== V7.2: SAFE JSON EXTRACTION ====================

def extract_and_validate_plan(raw_response: str) -> PlannerResponse:
    """
    V7.2: Index-based JSON extraction with PlannerFormatError on any failure.
    Refuses split() to avoid silent IndexError.
    """
    if "<JSON_START>" not in raw_response or "<JSON_END>" not in raw_response:
        raise PlannerFormatError("Missing explicit JSON boundaries <JSON_START>/<JSON_END>.")

    start = raw_response.index("<JSON_START>") + len("<JSON_START>")
    end = raw_response.index("<JSON_END>")
    raw_json = raw_response[start:end].strip()

    try:
        plan = PlannerResponse.model_validate_json(raw_json)

        # Guard: DOS Prevention
        if len(plan.requires_server_compute) > MAX_SERVER_STEPS:
            raise PlannerFormatError(
                f"DOS Guard: Plan requested {len(plan.requires_server_compute)} steps, max is {MAX_SERVER_STEPS}."
            )

        return plan
    except (ValidationError, json.JSONDecodeError) as e:
        raise PlannerFormatError(f"Validation failed: {e}")
    except PlannerFormatError:
        raise
    except Exception as e:
        raise PlannerFormatError(f"Unexpected extraction error: {e}")


# ==================== V7.2: PROPOSAL ENGINE ====================

class ProposalEngine:
    """
    V7.2 Deterministic Governed Reasoning — Pedagogical Planner (LLM #1)
    
    The LLM receives only AST metadata (IDs and operations) and returns
    a structured Plan (JSON). It NEVER writes math expressions — only references AST Node IDs.
    """

    def __init__(self, llm_gateway):
        self.llm_gateway = llm_gateway

    def _build_system_prompt(self, ast_metadata: dict, prompt_specialization: str, sub_question_text: str = "") -> str:
        """Build the strict Planner system prompt including boundary enforcement."""
        allowed_concepts = get_allowed_concepts(
            ast_metadata.get("detected_operations", ["general"])[0],
            "general_algebra"
        )
        concepts_str = ", ".join(allowed_concepts) if allowed_concepts else "מונחי אלגברה מדויקים"

        sub_q_line = f"\nSUB-QUESTION TO SOLVE NOW: {sub_question_text}" if sub_question_text else ""

        return f"""You are a Pedagogical Planner for a math tutoring system. 
Your role is to plan HOW to solve a math problem, NOT to solve it yourself.

PROBLEM CONTEXT (AST Metadata):
{json.dumps(ast_metadata, ensure_ascii=False, indent=2)}{sub_q_line}

CRITICAL RULES:
1. You MUST wrap your entire JSON response between <JSON_START> and <JSON_END> tags.
2. Do NOT write any math expressions or compute anything yourself.
3. When referring to AST nodes, use their IDs only (e.g. "ast_node_0").
4. Your response MUST conform to this exact JSON schema:
{{
    "pedagogical_rationale": "<short explanation of the teaching strategy in Hebrew>",
    "requires_server_compute": [
        {{"action": "<ENUM_ACTION>", "target_step_ref": "<ast_node_id>"}},
        ...
    ]
}}
5. Allowed actions (ENUM only): SOLVE_EQUATION, SIMPLIFY, FIND_DERIVATIVE, FIND_INTEGRAL, FACTOR, EXPAND, SUBSTITUTE, FIND_AXIS_INTERSECTIONS, CALCULATE_SLOPE_AND_LINE, CALCULATE_DISTANCE.
6. Maximum {MAX_SERVER_STEPS} compute steps.
7. Pedagogical Constraint: {prompt_specialization}
8. Vocabulary: Prioritize these terms: [{concepts_str}].
9. SINGLE-ACTION CONSTRAINT (critical): For complex equations (circles, trigonometry, systems),
   always use SOLVE_EQUATION directly on ast_node_0. DO NOT chain operations that reference
   intermediate result nodes (ast_node_1, ast_node_2, etc.) — those nodes do not exist in
   the server registry. The SymPy engine handles simplification and expansion internally
   as part of SOLVE_EQUATION. Violating this rule causes a server crash.
10. CONTEXTUAL ROUTING (V7.3 — critical): Analyze the SUB-QUESTION TO SOLVE NOW carefully.
    Select the action that matches the specific task:
    - "חיתוך", "צירים" → FIND_AXIS_INTERSECTIONS
    - "ישר", "שיפוע", "משואה" → CALCULATE_SLOPE_AND_LINE
    - "מרחק", "נקודה על המעגל" → CALCULATE_DISTANCE
    - General algebra/equation → SOLVE_EQUATION
    You MUST choose the action that solves specifically this sub-question.
    Do NOT use SOLVE_EQUATION when a more specific action applies."""


    async def generate_draft_proposal(self, context_obj, prompt_specialization: str, ast_metadata: dict = None) -> dict:
        """
        Asks LLM #1 (Planner) to produce a structured action plan.
        Returns {"success": True, "plan": PlannerResponse} or {"success": False, "reason": ...}.
        """
        if ast_metadata is None:
            ast_metadata = {}

        system_prompt = self._build_system_prompt(
            ast_metadata, prompt_specialization,
            sub_question_text=getattr(context_obj, 'sub_question_text', '')
        )
        user_prompt = (
            f"Plan the solution for: {context_obj.math_input}\n"
            f"Remember: Output JSON only, wrapped in <JSON_START>...<JSON_END> tags."
        )

        logger.info(f"[PLANNER] Requesting plan for: {context_obj.math_input}")

        for attempt in range(2):  # Max 1 retry (Full Strategy Re-run)
            try:
                raw_response = await self.llm_gateway.generate_raw(system_prompt, user_prompt)
                plan = extract_and_validate_plan(raw_response)

                logger.info(
                    f"[PLANNER] Valid plan received on attempt {attempt + 1}. "
                    f"Actions: {[t.action for t in plan.requires_server_compute]}"
                )
                # Phase 1 Live: Track which Enum actions the Planner chose (drift detection)
                chosen_actions = [task.action.value for task in plan.requires_server_compute]
                telemetry.emit_planner_strategy_distribution(chosen_actions)
                return {"success": True, "plan": plan}

            except PlannerFormatError as e:
                logger.warning(f"[PLANNER] Attempt {attempt + 1} failed: {e}")
                if attempt == 0:
                    # Full Strategy Re-run: inject hard feedback and retry
                    user_prompt = (
                        f"Your previous response was invalid: {e}\n"
                        f"Try again. Plan the solution for: {context_obj.math_input}\n"
                        f"You MUST use <JSON_START> and <JSON_END> tags. Output valid JSON only."
                    )
                    continue
                else:
                    logger.error("[PLANNER] Max retries exhausted. Failing closed.")
                    return {"success": False, "reason": str(e)}

            except Exception as e:
                logger.error(f"[PLANNER] LLM failure on attempt {attempt + 1}: {e}")
                if attempt == 0:
                    continue
                return {"success": False, "reason": f"LLM_FAILURE: {e}"}

        return {"success": False, "reason": "MAX_RETRIES_EXHAUSTED"}