agents/geometry_parser_agent.py

import json
import logging
import re
from typing import Dict, Any, Optional, Tuple
from dotenv import load_dotenv

load_dotenv()
logger = logging.getLogger(__name__)

from agents.runtime import get_agent_runtime, AgentRuntime


class GeometryParserAgent:
    """
    Unified Geometry Parser Agent (v7.0 - Agent Runtime & Cascading Controller):
    Directly extracts semantic entities, dimensions, target question,
    and generates high-precision Geometry DSL in a single, high-fidelity LLM inference step.
    """

    def __init__(self, runtime: Optional[AgentRuntime] = None):
        self.runtime = runtime or get_agent_runtime()

    def _clean_json(self, raw: str) -> str:
        s = raw.strip()
        json_match = re.search(r"```(?:json)?(.*?)```", s, re.DOTALL)
        if json_match:
            return json_match.group(1).strip()
        brace_match = re.search(r"(\{.*\})", s, re.DOTALL)
        if brace_match:
            return brace_match.group(1).strip()
        return s.strip()

    def _validate_parser_output(self, raw: str) -> Tuple[bool, Any]:
        """Validates JSON structure and extracts DSL."""
        try:
            cleaned = self._clean_json(raw)
            data = json.loads(cleaned)
            if not isinstance(data, dict):
                return False, "Output must be a JSON object"
            if "type" not in data and "geometry_dsl" not in data:
                return False, "Missing required 'type' or 'geometry_dsl' fields"
            dsl = data.get("geometry_dsl", "")
            if isinstance(dsl, list):
                dsl = "\n".join(dsl)
            data["geometry_dsl"] = dsl.strip()
            return True, data
        except Exception as e:
            return False, f"JSON parse error: {e}"

    async def process(
        self,
        text: str,
        feedback: Optional[str] = None,
        context: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        logger.info(f"==[GeometryParserAgent] Parsing problem & generating DSL (len={len(text)}) (v7.0)==")
        if feedback:
            logger.warning(f"[GeometryParserAgent] Feedback from previous attempt: {feedback}")
        if context:
            logger.info(f"[GeometryParserAgent] Using previous context (dsl_len={len(context.get('geometry_dsl', ''))})")

        system_prompt = """You are an expert Geometry Parser & DSL Generator (v6.0 - Hybrid Procedural & 3D Analytical).
Analyze the Vietnamese/LaTeX mathematical geometry problem and extract both the structured semantics AND the executable Geometry DSL program in a single step.

=== DSL SPECIFICATION v6.0 ===
-- 2D & 3D Basic Primitives --
POINT(A)                             — declare a point (supports A, B, A1, B1, A', B', S, O, M, N, H, I, G, D)
POINT(A, x, y, z)                    — declare a point with explicit coordinates
LENGTH(AB, 5)                        — distance between A and B is 5
ANGLE(A, 90)                         — angle at vertex A is 90° (or ANGLE(A, B, C, 60))
PARALLEL(AB, CD)                     — segment AB is parallel to CD
PERPENDICULAR(AB, CD)                — segment AB is perpendicular to CD
MIDPOINT(M, AB)                      — M is the midpoint of segment AB
SECTION(E, A, C, k)                  — E satisfies vector AE = k * vector AC (k is decimal, e.g. 0.5)
LINE(A, B)                           — infinite line passing through A and B
RAY(A, B)                            — ray starting at A and passing through B
CIRCLE(O, 5)                         — circle with center O and radius 5
SEGMENT(M, N)                        — auxiliary segment MN to be drawn
POLYGON_ORDER(A, B, C, D)            — polygon boundary vertex ordering
TRIANGLE(ABC)                        — triangle
SQUARE(ABCD)                         — square
RECTANGLE(ABCD)                      — rectangle
PARALLELOGRAM(ABCD)                  — parallelogram
RHOMBUS(ABCD)                        — rhombus
TRAPEZOID(ABCD)                      — trapezoid

-- Deterministic Construction Primitives (O(1) analytic, GeoBuildBench-inspired) --
INTERSECT_LINES(P, AB, CD)           — P is intersection of lines AB and CD
INTERSECT_LINE_PLANE(P, Line, Plane) — P is intersection of line and plane (e.g. INTERSECT_LINE_PLANE(M, SO, ABCD))
INCENTER(I, A, B, C)                 — I is incenter of triangle ABC (tâm đường tròn nội tiếp)
CIRCUMCENTER(O, A, B, C)             — O is circumcenter of triangle ABC (tâm đường tròn ngoại tiếp)
ORTHOCENTER(H, A, B, C)              — H is orthocenter of triangle ABC (trực tâm)
CENTROID(G, A, B, C)                 — G is centroid of triangle ABC (trọng tâm)
ANGLE_BISECTOR(D, A, B, C)           — D is foot of internal angle bisector of angle BAC on side BC
PROJECT_POINT_PLANE(H, P, ABC)       — H is orthogonal projection of point P onto plane ABC (H in plane, PH ⊥ plane)
PROJECT_POINT_LINE(H, P, AB)         — H is orthogonal projection of point P onto line AB

-- Micro-Predicates & Topological Constraints (MagicGeo-inspired) --
ANGLE_RELATION(A, B, C, D, E, F, 2.0)— angle(ABC) = 2.0 * angle(DEF)
ANGLE_SUM(A, B, C, D, E, F, 90)      — angle(ABC) + angle(DEF) = 90°
INSIDE_SEGMENT(M, A, B)              — point M lies strictly inside segment AB (0 < t < 1)
ON_RAY(M, A, B)                      — point M lies on ray AB extended past B (t > 1)
RATIO_SEGMENT(A, M, B, 2.0)          — point M divides segment with length ratio MA / MB = 2.0

-- 3D Polyhedrons & Round Solids --
PYRAMID(S_ABCD)                      — pyramid with apex S and base ABCD (supports S_ABC, S_ABCD)
PRISM(ABC_DEF)                       — triangular prism
PRISM(ABCD_A1B1C1D1)                 — quadrilateral prism
TETRAHEDRON(ABCD)                    — tetrahedron
CUBE(ABCD_A1B1C1D1)                  — cube
CUBOID(ABCD_A1B1C1D1)                — rectangular cuboid
FRUSTUM(ABCD_A1B1C1D1)               — frustum of a pyramid
CONE(S_O, r, h)                      — cone with apex S, base center O, radius r, height h
CYLINDER(O1_O2, r)                   — cylinder with axis O1-O2, radius r
SPHERE(O, r)                         — sphere with center O and radius r

-- 3D Spatial Relations & Oxyz Analytical Geometry --
PERPENDICULAR_PLANE(SA, ABCD)        — line SA ⊥ plane ABCD
ANGLE_LINE_PLANE(Line, Plane, 45)    — angle between line and plane is 45°
DIHEDRAL_ANGLE(Plane1, Plane2, 60)   — dihedral angle between two planes is 60°
DISTANCE_SKEW_LINES(AB, CD, 5)       — distance between skew lines AB and CD is 5
DISTANCE_POINT_PLANE(P, ABC, 3)      — distance from point P to plane ABC is 3
PLANE_EQ(P, A, B, C, D)              — plane P: Ax + By + Cz + D = 0
SPHERE_EQ(S, center, R)              — sphere S with center and radius R
TANGENT_PLANE_SPHERE(P, S)           — plane P is tangent to sphere S

=== OUTPUT FORMAT ===
Output ONLY a JSON object with this EXACT structure (no markdown, no extra keys):
{
    "type": "cube|cuboid|tetrahedron|cone|cylinder|frustum|pyramid|prism|sphere|rectangle|triangle|circle|parallelogram|trapezoid|square|rhombus|oxyz|general",
    "entities": ["Point S", "Point A", "Point B", "Point C", "Point D", "Point O"],
    "values": {"AB": 10, "SO": 15},
    "target_question": "Tính thể tích khối chóp S.ABCD",
    "analysis": "Tóm tắt bài toán ngắn gọn bằng tiếng Việt.",
    "geometry_dsl": "PYRAMID(S_ABCD)\\nSQUARE(ABCD)\\nLENGTH(AB, 4)\\nLENGTH(SA, 5)\\nPERPENDICULAR_PLANE(SA, ABCD)"
}

=== RULES ===
1. For pyramids with explicit base and height:
   PYRAMID(S_ABCD)
   SQUARE(ABCD)
   LENGTH(AB, 4)
   PERPENDICULAR_PLANE(SA, ABCD)
2. For projection/height foot: use PROJECT_POINT_PLANE(H, S, ABCD) or PERPENDICULAR_PLANE.
3. For Oxyz problems with given coordinates: declare POINT(A, x, y, z) directly.
4. Keep DSL commands clean, uppercase, and syntactically valid.
"""

        user_content = f"Đề bài toán:\n{text}"
        if context:
            user_content = f"PREVIOUS CONTEXT:\n{context.get('analysis', '')}\nDSL:\n{context.get('geometry_dsl', '')}\n\nNEW REQUEST:\n{text}"

        if feedback:
            user_content += f"\n\nPhản hồi từ lần chạy trước: {feedback}. Vui lòng sửa lại DSL và ràng buộc chính xác."

        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_content},
        ]

        try:
            data = await self.runtime.run(
                agent="geometry_parser",
                messages=messages,
                validator=self._validate_parser_output,
            )
        except Exception as e:
            logger.warning(f"[GeometryParserAgent] Agent runtime cascade failed: {e}. Using fallback structure.")
            data = {
                "type": "general",
                "entities": [],
                "values": {},
                "target_question": text,
                "analysis": text,
                "geometry_dsl": "",
            }

        dsl = data.get("geometry_dsl", "")
        if isinstance(dsl, list):
            dsl = "\n".join(dsl)
        data["geometry_dsl"] = dsl.strip()

        logger.info(f"[GeometryParserAgent] Success: type={data.get('type')}, dsl_lines={len(data['geometry_dsl'].splitlines())}")
        return data