pcb-afm-simulation-system/web/backend/app/services/plan_generator.py

"""AI Plan Generator (P3-M3).

Converts natural language requirements into structured simulation plans
using Kimi k3 model, with L0 pre-screening validation.
"""
import json
from pathlib import Path
from typing import Dict, List, Optional, Any, Tuple

from ..config import PROMPTS_DIR
from ..services.ai_client import get_kimi_client
from ..services.l0_prescreening import L0PreScreeningEngine


class AIPlanGenerator:
    """Generate simulation plans from natural language using AI."""

    def __init__(self, l0_engine: Optional[L0PreScreeningEngine] = None):
        self.ai_client = get_kimi_client()
        self.l0_engine = l0_engine or L0PreScreeningEngine()
        self._prompt_template = None

    def _load_prompt(self) -> str:
        """Load plan generation prompt template."""
        if self._prompt_template is None:
            prompt_path = PROMPTS_DIR / "plan_generation" / "generate.txt"
            if prompt_path.exists():
                with open(prompt_path, "r", encoding="utf-8") as f:
                    self._prompt_template = f.read()
            else:
                self._prompt_template = self._default_prompt()
        return self._prompt_template

    def _default_prompt(self) -> str:
        """Fallback default prompt."""
        return """\u4f60\u662f\u8f74\u5411\u78c1\u901a\u7535\u673a\u4eff\u771f\u65b9\u6848\u4e13\u5bb6\u3002\u5c06\u7528\u6237\u9700\u6c42\u8f6c\u5316\u4e3aJSON\u683c\u5f0f\u7684\u4eff\u771f\u65b9\u6848\uff0c\u5305\u542bplan_name\u3001topology\u3001boundary_conditions\u3001scan_variables\u3001search_strategy\u3001acceptance_criteria\u3001reasoning\u3002\u8f93\u51fa\u7eafJSON\u3002"""

    def generate(
        self,
        user_requirement: str,
        project_context: Optional[Dict[str, Any]] = None,
        existing_experience: Optional[List[Dict[str, Any]]] = None,
    ) -> Dict[str, Any]:
        """Generate a simulation plan from natural language.

        Args:
            user_requirement: Natural language description of the simulation need.
            project_context: Optional project context (topology, existing boundary conditions).
            existing_experience: Optional similar experience cases for reference.

        Returns:
            Dict with generated plan, validation results, and AI reasoning.
        """
        if not self.ai_client.is_configured:
            raise RuntimeError("KIMI_API_KEY is not configured")

        # Build context message
        context_parts = []
        if project_context:
            context_parts.append(f"\u9879\u76ee\u4e0a\u4e0b\u6587\uff1a{json.dumps(project_context, ensure_ascii=False)}")
        if existing_experience:
            context_parts.append(f"\u53c2\u8003\u7ecf\u9a8c\u6848\u4f8b\uff08{len(existing_experience)}\u4e2a\uff09\uff1a{json.dumps(existing_experience[:3], ensure_ascii=False)}")

        user_message = user_requirement
        if context_parts:
            user_message += "\n\n" + "\n".join(context_parts)

        # Call AI
        system_prompt = self._load_prompt()
        result = self.ai_client.chat_json(
            messages=[{"role": "user", "content": user_message}],
            system_prompt=system_prompt,
            max_tokens=2000,
        )

        # Parse generated plan
        plan = result.get("parsed_json", {})
        if not plan:
            # Try to extract and repair JSON from raw content
            raw = result.get("raw_content", result.get("content", ""))
            plan = self._try_repair_json(raw)
            if not plan:
                plan = {"error": "Failed to parse AI response", "raw_content": raw[:1000]}

        # Validate with L0 pre-screening
        validation = self._validate_plan(plan)

        return {
            "plan": plan,
            "validation": validation,
            "ai_reasoning": plan.get("reasoning", ""),
            "usage": result.get("usage", {}),
            "raw_content": result.get("content", ""),
        }

    def _try_repair_json(self, raw: str) -> Dict[str, Any]:
        """Try to extract and repair JSON from potentially truncated AI output.

        Handles:
        - Markdown code fences
        - Truncated JSON (missing closing braces)
        - Extra text before/after JSON
        """
        if not raw:
            return {}

        text = raw.strip()

        # Remove markdown code fences
        if text.startswith("```"):
            lines = text.split("\n")
            if lines[0].startswith("```"):
                lines = lines[1:]
            if lines and lines[-1].strip() == "```":
                lines = lines[:-1]
            text = "\n".join(lines).strip()

        # Find first { and last }
        first_brace = text.find("{")
        last_brace = text.rfind("}")

        if first_brace == -1:
            return {}

        if last_brace == -1 or last_brace < first_brace:
            # JSON is truncated - try to close it
            json_str = text[first_brace:]
            # Count open braces and close them
            open_braces = json_str.count("{") - json_str.count("}")
            open_brackets = json_str.count("[") - json_str.count("]")
            # Remove trailing incomplete content
            last_comma = max(json_str.rfind(","), json_str.rfind(":"))
            if last_comma > len(json_str) * 0.8:
                json_str = json_str[:last_comma]
            # Close brackets and braces
            json_str += "]" * max(0, open_brackets)
            json_str += "}" * max(0, open_braces)
        else:
            json_str = text[first_brace:last_brace + 1]

        try:
            return json.loads(json_str)
        except json.JSONDecodeError:
            # Try one more repair: remove trailing comma before closing
            try:
                repaired = json_str.rstrip()
                while repaired and repaired[-1] in ", \n\t":
                    repaired = repaired[:-1]
                # Re-count and close
                open_braces = repaired.count("{") - repaired.count("}")
                open_brackets = repaired.count("[") - repaired.count("]")
                repaired += "]" * max(0, open_brackets)
                repaired += "}" * max(0, open_braces)
                return json.loads(repaired)
            except (json.JSONDecodeError, Exception):
                return {}

    def _validate_plan(self, plan: Dict[str, Any]) -> Dict[str, Any]:
        """Validate generated plan with L0 pre-screening.

        Checks:
        1. Boundary conditions feasibility
        2. Scan variable ranges within engineering limits
        3. Sample point count estimation
        4. Required fields presence
        """
        issues = []
        warnings = []
        checks = []

        # Check required fields
        required_fields = ["plan_name", "topology", "scan_variables"]
        for field in required_fields:
            if field not in plan:
                issues.append(f"Missing required field: {field}")

        # Validate topology
        topology = plan.get("topology", "")
        if topology and topology not in ("SSSR", "DRSS", "SDSR"):
            warnings.append(f"Unusual topology: {topology}")

        # Validate boundary conditions with L0
        bc = plan.get("boundary_conditions", {})
        if bc:
            l0_report = self.l0_engine.evaluate(bc)
            checks.append({
                "name": "boundary_conditions_l0",
                "feasible": l0_report.feasible,
                "passed": l0_report.passed_checks,
                "total": l0_report.total_checks,
                "failed_items": [r.name for r in l0_report.results if not r.passed],
            })
            if not l0_report.feasible:
                issues.append(f"Boundary conditions infeasible: {l0_report.failed_checks} checks failed")

        # Validate scan variables
        variables = plan.get("scan_variables", [])
        if not isinstance(variables, list):
            variables = []
        total_points = 1
        normalized_vars = []
        for var in variables:
            if isinstance(var, str):
                var = {"name": var, "min_value": None, "max_value": None}
            if not isinstance(var, dict):
                continue
            normalized_vars.append(var)
            name = var.get("name", "unknown")
            min_v = var.get("min_value")
            max_v = var.get("max_value")
            step = var.get("step")

            if min_v is None or max_v is None:
                issues.append(f"Variable {name}: missing min/max values")
                continue

            if min_v >= max_v:
                issues.append(f"Variable {name}: min ({min_v}) >= max ({max_v})")

            if step and step > 0:
                n_points = int((max_v - min_v) / step) + 1
                total_points *= n_points
                if n_points > 50:
                    warnings.append(f"Variable {name}: {n_points} levels may be too many")
        plan["scan_variables"] = normalized_vars

        checks.append({
            "name": "scan_variables",
            "count": len(variables),
            "estimated_full_factorial_points": total_points,
            "adaptive_search_recommended": total_points > 100,
        })

        if total_points > 500:
            warnings.append(f"Full factorial would require {total_points} points - strongly recommend adaptive search")

        # Validate search strategy
        search = plan.get("search_strategy", {})
        if search:
            max_calls = search.get("max_solver_calls", 80)
            if max_calls < total_points and not search.get("method", "").startswith(("constrained", "active")):
                warnings.append(f"Budget ({max_calls}) < full factorial ({total_points}) but method is not adaptive")

        overall_valid = len(issues) == 0
        return {
            "valid": overall_valid,
            "issues": issues,
            "warnings": warnings,
            "checks": checks,
        }

    def refine_plan(
        self,
        original_plan: Dict[str, Any],
        user_feedback: str,
    ) -> Dict[str, Any]:
        """Refine an existing plan based on user feedback.

        Args:
            original_plan: The previously generated plan.
            user_feedback: Natural language feedback for refinement.

        Returns:
            Refined plan with validation.
        """
        if not self.ai_client.is_configured:
            raise RuntimeError("KIMI_API_KEY is not configured")

        system_prompt = self._load_prompt() + "\n\n\u4f60\u6b63\u5728\u4f18\u5316\u4e00\u4e2a\u5df2\u6709\u7684\u4eff\u771f\u65b9\u6848\u3002\u6839\u636e\u7528\u6237\u53cd\u9988\u8c03\u6574\u65b9\u6848\uff0c\u4fdd\u6301\u5176\u4ed6\u90e8\u5206\u4e0d\u53d8\u3002"

        user_message = f"\u539f\u6709\u65b9\u6848\uff1a\n{json.dumps(original_plan, ensure_ascii=False, indent=2)}\n\n\u7528\u6237\u53cd\u9988\uff1a{user_feedback}\n\n\u8bf7\u8f93\u51fa\u4f18\u5316\u540e\u7684\u5b8c\u6574\u65b9\u6848JSON\u3002"

        result = self.ai_client.chat_json(
            messages=[{"role": "user", "content": user_message}],
            system_prompt=system_prompt,
            max_tokens=2000,
        )

        refined_plan = result.get("parsed_json", {})
        if not refined_plan:
            raw = result.get("raw_content", result.get("content", ""))
            refined_plan = self._try_repair_json(raw)
            if not refined_plan:
                refined_plan = {"error": "Failed to parse AI response", "raw_content": raw[:1000]}

        validation = self._validate_plan(refined_plan)

        return {
            "plan": refined_plan,
            "validation": validation,
            "ai_reasoning": refined_plan.get("reasoning", ""),
            "usage": result.get("usage", {}),
        }


# Global singleton
_generator: Optional[AIPlanGenerator] = None


def get_plan_generator() -> AIPlanGenerator:
    """Get or create global AIPlanGenerator singleton."""
    global _generator
    if _generator is None:
        _generator = AIPlanGenerator()
    return _generator