pcb-afm-simulation-system/web/backend/app/routers/ai_plan.py

"""AI Plan Generation router (P3-M3).

Endpoints for generating simulation plans from natural language
using Kimi k3 model, with L0 pre-screening validation.

Includes project-bound endpoint that saves AI-generated plan directly
as a SimulationPlan under the project.
"""
import json
import uuid
from datetime import datetime
from typing import Dict, Any, Optional, List
from fastapi import APIRouter, HTTPException, Depends
from pydantic import BaseModel, Field
from sqlalchemy.orm import Session

from ..database import get_db
from ..models.project import Project
from ..models.simulation_plan import SimulationPlan
from ..models.experience_case import ExperienceCase
from ..services.plan_generator import get_plan_generator

router = APIRouter(prefix="/api/ai-plan", tags=["AI Plan Generation"])


def _load_experience_cases(db: Session, topology: str, limit: int = 5) -> List[Dict[str, Any]]:
    """Load recent experience cases for the topology so the AI can reuse prior
    design knowledge. Without this the library accumulated by the adaptive
    loop never flows back into plan generation (value-chain break)."""
    cases = (
        db.query(ExperienceCase)
        .filter(ExperienceCase.topology == topology)
        .order_by(ExperienceCase.created_at.desc())
        .limit(limit)
        .all()
    )
    out = []
    for c in cases:
        out.append({
            "topology": c.topology,
            "params": c.get_params(),
            "metrics": c.get_metrics(),
            "conclusion": c.conclusion,
            "tags": c.tags,
        })
    return out


class GeneratePlanRequest(BaseModel):
    """Request to generate a plan from natural language."""
    requirement: str = Field(..., description="Natural language simulation requirement")
    project_context: Optional[Dict[str, Any]] = Field(default=None, description="Optional project context")
    existing_experience: Optional[List[Dict[str, Any]]] = Field(default=None, description="Optional similar experience cases")


class GenerateAndSaveRequest(BaseModel):
    """Request to generate a plan and save it under a project."""
    requirement: str = Field(default="", description="Natural language requirement (optional, uses project BC as base)")
    model_path: Optional[str] = Field(default=None, description="Override model path")
    plan_name: Optional[str] = Field(default=None, description="Override plan name")


class RefinePlanRequest(BaseModel):
    """Request to refine an existing plan."""
    original_plan: Dict[str, Any] = Field(..., description="The previously generated plan")
    feedback: str = Field(..., description="Natural language feedback for refinement")


@router.post("/generate")
def generate_plan(request: GeneratePlanRequest, db: Session = Depends(get_db)):
    """Generate a simulation plan from natural language.

    Uses Kimi k3 model to convert natural language requirements into
    a structured simulation plan, with L0 pre-screening validation.
    Returns unified plan format (fixed_params + variables with values).
    """
    try:
        generator = get_plan_generator()
        # Auto-load library experience when the caller didn't pass any.
        experience = request.existing_experience
        if experience is None:
            topo = (request.project_context or {}).get("topology", "SSSR")
            experience = _load_experience_cases(db, topo)
        result = generator.generate(
            user_requirement=request.requirement,
            project_context=request.project_context,
            existing_experience=experience,
        )
        return result
    except RuntimeError as e:
        raise HTTPException(status_code=503, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Plan generation failed: {str(e)}")


@router.post("/projects/{project_id}/generate-and-save")
def generate_and_save_plan(
    project_id: int,
    request: GenerateAndSaveRequest,
    db: Session = Depends(get_db),
):
    """Generate a plan using AI and save it directly under a project.

    Uses the project's boundary conditions as context. The generated plan
    is in unified format (fixed_params + variables with value lists) and
    is saved as a SimulationPlan with status 'draft'.

    Args:
        project_id: Target project ID.
        request: Optional natural language requirement and overrides.

    Returns:
        Created SimulationPlan response dict.
    """
    project = db.query(Project).filter(Project.id == project_id).first()
    if not project:
        raise HTTPException(status_code=404, detail="Project not found")

    # Build context from project boundary conditions. Expand against the full
    # BC catalog so unset fields appear explicitly as null - the AI can only
    # propose bc_suggestions for fields it can SEE are unset (previously unset
    # keys were simply absent, so the model had no way to know what to fill).
    from ..services.bc_fields import get_bc_field_catalog
    bc = project.get_boundary_conditions()
    full_bc = {f["key"]: bc.get(f["key"]) for f in get_bc_field_catalog()}
    context = {
        "project_id": project.id,
        "project_name": project.name,
        "topology": project.topology,
        "model_path": project.model_path or "",
        "boundary_conditions": full_bc,
    }

    # Build requirement: use explicit requirement or construct from BC
    if request.requirement.strip():
        requirement = request.requirement
    else:
        # Construct a requirement from boundary conditions
        parts = [f"Topology: {project.topology}"]
        if bc.get("outer_diameter_mm"):
            parts.append(f"Outer diameter <= {bc['outer_diameter_mm']}mm")
        if bc.get("speed_rpm"):
            parts.append(f"Rated speed {bc['speed_rpm']}rpm")
        if bc.get("current_a"):
            parts.append(f"RMS current {bc['current_a']}A")
        if bc.get("target_efficiency_pct"):
            parts.append(f"Target efficiency >= {bc['target_efficiency_pct']}%")
        if bc.get("target_torque_nm"):
            parts.append(f"Target torque >= {bc['target_torque_nm']}Nm")
        if bc.get("max_losses_w"):
            parts.append(f"Max losses <= {bc['max_losses_w']}W")
        if bc.get("power_w"):
            parts.append(f"Power {bc['power_w']}W")
        requirement = "Generate a simulation plan for: " + "; ".join(parts) + ". Recommend scan variables and fixed parameters."

    try:
        generator = get_plan_generator()
        result = generator.generate(
            user_requirement=requirement,
            project_context=context,
            existing_experience=_load_experience_cases(db, project.topology or "SSSR"),
        )
    except RuntimeError as e:
        raise HTTPException(status_code=503, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"AI generation failed: {str(e)}")

    unified_plan = result.get("plan", {})
    if not unified_plan or "variables" not in unified_plan:
        raise HTTPException(status_code=502, detail="AI returned invalid plan format")

    # Override model_path and name
    model_path = request.model_path or project.model_path or unified_plan.get("model_path", "")
    plan_name = request.plan_name or unified_plan.get("name", f"AI Plan {datetime.now().strftime('%m%d %H%M')}")

    # Normalize topology against the single-source registry. Legacy projects may
    # carry a non-registered topology (e.g. "AFIR"); fall back to the only active
    # topology (SSSR) and surface a warning instead of failing with 422.
    from src.afmcore.topology import is_supported as _topology_supported
    topology = project.topology or "SSSR"
    topology_warnings: List[str] = []
    if not _topology_supported(topology):
        topology_warnings.append(
            f"Unknown topology '{topology}' - defaulted to SSSR"
        )
        topology = "SSSR"

    # Auto-fill the base model from the topology registry when nothing
    # provided one (project/model-less plans otherwise fail preflight with an
    # empty model_path). DRSS/SDSR have no base model yet -> stays empty and
    # preflight will report it explicitly.
    if not model_path:
        from src.afmcore.topology import default_model_for
        model_path = default_model_for(topology)
        if model_path:
            topology_warnings.append(
                f"model_path auto-filled from topology default model: {model_path}"
            )

    # Build plan_data
    plan_id = f"SP-{datetime.now().strftime('%Y%m%d-%H%M%S')}-{uuid.uuid4().hex[:6]}"
    # Persist the project's boundary conditions (normalized to canonical BC
    # keys) onto the plan so the plan-detail boundary display has data and the
    # BC the plan was generated against is traceable (P1-1).
    # Source tracking (bc_meta): user-specified project BC keys are tagged
    # "user"; AI bc_suggestions only fill keys the user left unset and are
    # tagged "ai" (requirement: distinguish user-specified vs AI-supplemented).
    from ..services.bc_fields import normalize_bc
    user_bc = normalize_bc(bc)
    raw_ai_plan = result.get("raw_ai_plan", {})
    ai_suggestions = raw_ai_plan.get("bc_suggestions") if isinstance(raw_ai_plan, dict) else None
    bc_meta = {k: {"source": "user"} for k in user_bc}
    merged_bc = dict(user_bc)
    if isinstance(ai_suggestions, dict):
        for k, v in ai_suggestions.items():
            if v is None or v == "":
                continue
            norm = normalize_bc({k: v})
            for nk, nv in norm.items():
                if nk not in merged_bc:  # never override a user-specified value
                    merged_bc[nk] = nv
                    bc_meta[nk] = {"source": "ai"}
    plan_data = {
        "plan_id": plan_id,
        "plan_version": "2.0",
        "created_at": datetime.now().isoformat(timespec="seconds"),
        "topology": topology,
        "model_path": model_path,
        "boundary_conditions": merged_bc,
        "bc_meta": bc_meta,
        "fixed_params": unified_plan.get("fixed_params", []),
        "variables": unified_plan.get("variables", []),
        "cases": unified_plan.get("cases", [{"id": "default", "name": "Default", "params": {}}]),
        "output_metrics": unified_plan.get("output_metrics", []),
        "search_strategy": unified_plan.get("search_strategy"),
        "acceptance_criteria": unified_plan.get("acceptance_criteria"),
        "ai_reasoning": unified_plan.get("ai_reasoning", ""),
        "iteration": 1,
        "parent_plan_id": "",
    }

    # Validate against the single-source plan schema (structural checks
    # before persisting an AI-generated plan).
    from src.plan_schema import validate_plan_dict
    _ok, _errs = validate_plan_dict(plan_data, require_model_path=False)
    if not _ok:
        raise HTTPException(
            status_code=422,
            detail="Invalid generated plan: " + "; ".join(_errs),
        )

    # Compute estimated points
    estimated_points = 1
    for v in plan_data["variables"]:
        estimated_points *= len(v.get("values", [])) if v.get("values") else 1

    # Variables summary for display
    variables_summary = {}
    for v in plan_data["variables"]:
        variables_summary[v["name"]] = {
            "unit": v.get("unit", ""),
            "values": v.get("values", []),
            "count": len(v.get("values", [])),
        }

    # Save to database
    db_plan = SimulationPlan(
        project_id=project_id,
        name=plan_name,
        plan_id=plan_id,
        status="draft",
        estimated_points=estimated_points,
        estimated_time_min=estimated_points * 3,
        notes=f"AI-generated. Reasoning: {unified_plan.get('ai_reasoning', '')[:500]}",
    )
    db_plan.set_plan_dict(plan_data)
    db_plan.variables_summary = json.dumps(variables_summary, ensure_ascii=False)
    db.add(db_plan)
    db.commit()
    db.refresh(db_plan)

    return {
        "id": db_plan.id,
        "plan_id": db_plan.plan_id,
        "name": db_plan.name,
        "status": db_plan.status,
        "estimated_points": db_plan.estimated_points,
        "estimated_time_min": db_plan.estimated_time_min,
        "plan_data": plan_data,
        "ai_reasoning": unified_plan.get("ai_reasoning", ""),
        "validation": result.get("validation", {}),
        "warnings": topology_warnings + unified_plan.get("warnings", []),
    }


@router.post("/refine")
def refine_plan(request: RefinePlanRequest):
    """Refine an existing plan based on user feedback."""
    try:
        generator = get_plan_generator()
        result = generator.refine_plan(
            original_plan=request.original_plan,
            user_feedback=request.feedback,
        )
        return result
    except RuntimeError as e:
        raise HTTPException(status_code=503, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Plan refinement failed: {str(e)}")


@router.get("/templates")
def list_templates():
    """List available plan generation prompt templates."""
    return {
        "templates": [
            {
                "name": "generate",
                "path": "prompts/plan_generation/generate.txt",
                "description": "Natural language to structured simulation plan (unified v2.0)",
            }
        ],
        "model": "kimi-k3",
        "endpoint": "https://api.kimi.com/coding/v1",
    }