src/form_generator_workflow.py

"""
LangGraph workflow for form generation and refine.
Mirrors patterns from resume_workflow and visualization_workflow.
"""
from typing import TypedDict, List, Dict, Any, Optional, Literal
import json
import logging
import time
import traceback
import os

logger = logging.getLogger(__name__)

from langgraph.graph import StateGraph, END
from langchain_openai import ChatOpenAI

from prompts import GENERATE_SYSTEM, REFINE_SYSTEM, PLAN_QUESTIONS_SYSTEM
from openai_client import extract_json_from_response
from .docs_client import fetch_docs
from .context_pack_builder import detect_intent, build_context_packs


class PlanState(TypedDict):
    """State for clarification planning (0-5 questions)."""
    description: Optional[str]
    current_fields: Optional[List[Dict[str, Any]]]
    current_title: Optional[str]
    should_ask_questions: bool
    questions: List[Dict[str, Any]]
    reasoning_summary: Optional[str]
    error: Optional[str]
    success: bool


class FormGeneratorState(TypedDict, total=False):
    """State for form generator workflow."""
    # Mode: which path to take
    mode: Literal["generate", "refine"]
    # Inputs for generate
    description: Optional[str]
    # Inputs for refine
    user_request: Optional[str]
    current_fields: Optional[List[Dict[str, Any]]]
    current_title: Optional[str]
    # Optional context hints (from chat)
    conversation_context: Optional[List[str]]
    user_goals: Optional[List[str]]
    preferred_field_types: Optional[List[str]]
    must_have_logic: Optional[bool]
    # Results
    form_json: Optional[Dict[str, Any]]
    diff: Optional[Dict[str, Any]]
    metadata: Optional[Dict[str, Any]]
    # Status
    error: Optional[str]
    success: bool
    processing_time: float
    usage: Optional[Dict[str, int]]
    processing_complete: bool


class FormGeneratorWorkflow:
    """LangGraph workflow for form generation and surgical refine."""

    def __init__(self):
        self.workflow = self._build_workflow()

    def _build_workflow(self) -> StateGraph:
        workflow = StateGraph(FormGeneratorState)

        workflow.add_node("validate_input", self.validate_input_node)
        workflow.add_node("generate_form_node", self.generate_form_node)
        workflow.add_node("refine_form_node", self.refine_form_node)
        workflow.add_node("finalize_results", self.finalize_results_node)

        workflow.add_conditional_edges(
            "validate_input",
            self._route_by_mode,
            {"generate": "generate_form_node", "refine": "refine_form_node", "error": "finalize_results"},
        )
        workflow.add_edge("generate_form_node", "finalize_results")
        workflow.add_edge("refine_form_node", "finalize_results")
        workflow.add_edge("finalize_results", END)
        workflow.set_entry_point("validate_input")

        return workflow.compile()

    def _route_by_mode(self, state: FormGeneratorState) -> str:
        if state.get("error"):
            return "error"
        return state.get("mode", "error")

    def validate_input_node(self, state: FormGeneratorState) -> FormGeneratorState:
        """Validate input and set mode (generate vs refine)."""
        try:
            if state.get("description") and str(state["description"]).strip():
                state["mode"] = "generate"
                state["success"] = True
                return state
            if (
                state.get("user_request")
                and str(state["user_request"]).strip()
                and state.get("current_fields") is not None
            ):
                state["mode"] = "refine"
                state["success"] = True
                return state
            state["error"] = "Missing input: provide description (generate) or user_request + current_fields (refine)"
            state["success"] = False
        except Exception as e:
            state["error"] = f"Validation error: {str(e)}"
            state["success"] = False
        return state

    def generate_form_node(self, state: FormGeneratorState) -> FormGeneratorState:
        """Generate full form JSON from description using LLM."""
        try:
            description = (state.get("description") or "").strip()
            user_content = f"USER REQUEST:\n{description}\n\nGenerate the form JSON only (no markdown)."
            docs = fetch_docs()
            if docs:
                intent = detect_intent(
                    description=description,
                    conversation_context=state.get("conversation_context"),
                )
                if state.get("must_have_logic"):
                    intent["needs_conditional_logic"] = True
                context_packs = build_context_packs(docs, intent, log_usage=True)
                if context_packs:
                    logger.info("generate_form using dynamic context intent=%s", intent)
                    user_content = (
                        f"USER REQUEST:\n{description}\n\n"
                        f"REFERENCE (use for field types, conditional logic, document-extraction):\n{context_packs}\n\n"
                        "Generate the form JSON only (no markdown)."
                    )
            else:
                logger.debug("generate_form using static prompt only (no docs)")
            client = ChatOpenAI(
                api_key=os.getenv("OPENAI_API_KEY"),
                model="gpt-4.1-nano",
                temperature=0.3,
                max_tokens=4000,
                model_kwargs={"response_format": {"type": "json_object"}},
            )
            start = time.time()
            response = client.invoke(
                [{"role": "system", "content": GENERATE_SYSTEM}, {"role": "user", "content": user_content}]
            )
            elapsed = time.time() - start
            raw = response.content if hasattr(response, "content") else str(response)
            json_str = extract_json_from_response(raw)
            data = json.loads(json_str)
            form_json, metadata = self._validate_form_json(data)
            state["form_json"] = form_json
            state["metadata"] = metadata
            state["processing_time"] = elapsed
            state["success"] = True
            if hasattr(response, "response_metadata") and response.response_metadata.get("usage"):
                usage = response.response_metadata["usage"]
                state["usage"] = {
                    "input_tokens": usage.get("input_tokens", 0),
                    "output_tokens": usage.get("output_tokens", 0),
                    "total_tokens": usage.get("total_tokens", 0),
                }
        except json.JSONDecodeError as e:
            state["error"] = f"Generated form JSON is invalid: {e}"
            state["success"] = False
        except Exception as e:
            state["error"] = str(e)
            state["success"] = False
            traceback.print_exc()
        return state

    def refine_form_node(self, state: FormGeneratorState) -> FormGeneratorState:
        """Produce surgical diff (changes, additions, removals) using LLM."""
        try:
            user_request = (state.get("user_request") or "").strip()
            current_fields = state.get("current_fields") or []
            current_title = state.get("current_title") or "Untitled"
            fields_json = json.dumps(current_fields, indent=2)
            max_fields_chars = 12000
            if len(fields_json) > max_fields_chars:
                fields_json = fields_json[: max_fields_chars - 3] + "..."
            user_content = (
                f"Current form title: {current_title}\n\n"
                f"Current fields (JSON):\n{fields_json}\n\n"
                f"User request: {user_request}\n\n"
                "Return ONLY the diff JSON (changes, additions, removals). No markdown."
            )
            docs = fetch_docs()
            if docs:
                intent = detect_intent(
                    user_request=user_request,
                    current_fields=current_fields,
                    conversation_context=state.get("conversation_context"),
                )
                context_packs = build_context_packs(docs, intent, max_total_chars=4000, log_usage=True)
                if context_packs:
                    logger.info("refine_form using dynamic context intent=%s", intent)
                    user_content = (
                        f"Current form title: {current_title}\n\n"
                        f"Current fields (JSON):\n{fields_json}\n\n"
                        f"REFERENCE (for new/changed fields):\n{context_packs}\n\n"
                        f"User request: {user_request}\n\n"
                        "Return ONLY the diff JSON (changes, additions, removals). No markdown."
                    )
            else:
                logger.debug("refine_form using static prompt only (no docs)")
            client = ChatOpenAI(
                api_key=os.getenv("OPENAI_API_KEY"),
                model="gpt-4.1-nano",
                temperature=0.3,
                max_tokens=4000,
                model_kwargs={"response_format": {"type": "json_object"}},
            )
            start = time.time()
            response = client.invoke(
                [{"role": "system", "content": REFINE_SYSTEM}, {"role": "user", "content": user_content}]
            )
            elapsed = time.time() - start
            raw = response.content if hasattr(response, "content") else str(response)
            json_str = extract_json_from_response(raw)
            data = json.loads(json_str)
            changes = data.get("changes")
            additions = data.get("additions")
            removals = data.get("removals")
            if not isinstance(changes, list):
                changes = []
            if not isinstance(additions, list):
                additions = []
            if not isinstance(removals, list):
                removals = []
            state["diff"] = {"changes": changes, "additions": additions, "removals": removals}
            state["processing_time"] = elapsed
            state["success"] = True
            if hasattr(response, "response_metadata") and response.response_metadata.get("usage"):
                usage = response.response_metadata["usage"]
                state["usage"] = {
                    "input_tokens": usage.get("input_tokens", 0),
                    "output_tokens": usage.get("output_tokens", 0),
                    "total_tokens": usage.get("total_tokens", 0),
                }
        except json.JSONDecodeError as e:
            state["error"] = f"Refine diff JSON is invalid: {e}"
            state["success"] = False
        except Exception as e:
            state["error"] = str(e)
            state["success"] = False
            traceback.print_exc()
        return state

    def finalize_results_node(self, state: FormGeneratorState) -> FormGeneratorState:
        """Mark processing complete."""
        state["processing_complete"] = True
        return state

    @staticmethod
    def _validate_form_json(data: Dict[str, Any]) -> tuple:
        """Ensure form has title and fields; return (form_json, metadata)."""
        if not isinstance(data, dict):
            raise ValueError("Form must be a JSON object")
        title = data.get("title") or "Untitled Form"
        fields = data.get("fields")
        if not isinstance(fields, list):
            raise ValueError("Form must have a 'fields' array")
        form_json = {"title": title, "fields": fields}
        if data.get("description") is not None:
            form_json["description"] = data["description"]
        field_types = list({f.get("type") for f in fields if isinstance(f, dict) and f.get("type")})
        metadata = {
            "fieldCount": len(fields),
            "fieldTypes": field_types,
            "hasConditionalLogic": any(
                (f.get("conditionalLogic") or {}).get("enabled") for f in fields if isinstance(f, dict)
            ),
            "hasAIFields": any(f.get("type") == "document-extraction" for f in fields if isinstance(f, dict)),
        }
        return form_json, metadata

    def generate_form(
        self,
        description: str,
        conversation_context: Optional[List[str]] = None,
        user_goals: Optional[List[str]] = None,
        preferred_field_types: Optional[List[str]] = None,
        must_have_logic: Optional[bool] = None,
    ) -> Dict[str, Any]:
        """Run workflow in generate mode and return API-shaped result."""
        initial = FormGeneratorState(
            mode="generate",
            description=description.strip(),
            user_request=None,
            current_fields=None,
            current_title=None,
            conversation_context=conversation_context,
            user_goals=user_goals,
            preferred_field_types=preferred_field_types,
            must_have_logic=must_have_logic,
            form_json=None,
            diff=None,
            metadata=None,
            error=None,
            success=False,
            processing_time=0.0,
            usage=None,
            processing_complete=False,
        )
        try:
            final = self.workflow.invoke(initial)
            return {
                "formJSON": final.get("form_json"),
                "metadata": final.get("metadata") or {},
                "warnings": None,
                "success": final.get("success", False),
                "error": final.get("error"),
                "processing_time": final.get("processing_time", 0.0),
                "usage": final.get("usage"),
            }
        except Exception as e:
            return {
                "formJSON": None,
                "metadata": {},
                "warnings": None,
                "success": False,
                "error": str(e),
                "processing_time": 0.0,
                "usage": None,
            }

    def refine_form(
        self,
        user_request: str,
        current_fields: List[Dict[str, Any]],
        current_title: str,
        conversation_context: Optional[List[str]] = None,
    ) -> Dict[str, Any]:
        """Run workflow in refine mode and return API-shaped result."""
        initial = FormGeneratorState(
            mode="refine",
            description=None,
            user_request=user_request.strip(),
            current_fields=current_fields,
            current_title=current_title or "",
            conversation_context=conversation_context,
            form_json=None,
            diff=None,
            metadata=None,
            error=None,
            success=False,
            processing_time=0.0,
            usage=None,
            processing_complete=False,
        )
        try:
            final = self.workflow.invoke(initial)
            return {
                "diff": final.get("diff") or {"changes": [], "additions": [], "removals": []},
                "success": final.get("success", False),
                "error": final.get("error"),
                "processing_time": final.get("processing_time", 0.0),
                "usage": final.get("usage"),
            }
        except Exception as e:
            return {
                "diff": {"changes": [], "additions": [], "removals": []},
                "success": False,
                "error": str(e),
                "processing_time": 0.0,
                "usage": None,
            }


def _build_plan_workflow() -> StateGraph:
    """Build small graph for plan-form-request (clarification questions)."""
    workflow = StateGraph(PlanState)
    workflow.add_node("plan_questions", _plan_questions_node)
    workflow.add_edge("plan_questions", END)
    workflow.set_entry_point("plan_questions")
    return workflow.compile()


def _plan_questions_node(state: PlanState) -> PlanState:
    """Single node: call LLM to decide if we need clarifying questions."""
    try:
        description = (state.get("description") or "").strip()
        current_fields = state.get("current_fields")
        current_title = state.get("current_title") or ""
        client = ChatOpenAI(
            api_key=os.getenv("OPENAI_API_KEY"),
            model="gpt-4.1-nano",
            temperature=0.2,
            max_tokens=1500,
            model_kwargs={"response_format": {"type": "json_object"}},
        )
        user_content = f"User request: {description}"
        if current_fields or current_title:
            user_content += f"\nCurrent form title: {current_title}\nCurrent fields (count): {len(current_fields or [])}"
        user_content += "\n\nReturn the JSON only (should_ask_questions, reasoning_summary, questions)."
        response = client.invoke(
            [{"role": "system", "content": PLAN_QUESTIONS_SYSTEM}, {"role": "user", "content": user_content}]
        )
        raw = response.content if hasattr(response, "content") else str(response)
        json_str = extract_json_from_response(raw)
        data = json.loads(json_str)
        state["should_ask_questions"] = bool(data.get("should_ask_questions", False))
        state["reasoning_summary"] = str(data.get("reasoning_summary", "")).strip() or "Ready."
        qs = data.get("questions")
        if isinstance(qs, list):
            state["questions"] = [
                {
                    "id": q.get("id", f"q{i}"),
                    "label": q.get("label", ""),
                    "type": q.get("type", "text"),
                    "required": bool(q.get("required", True)),
                    "options": q.get("options") if isinstance(q.get("options"), list) else None,
                }
                for i, q in enumerate(qs) if isinstance(q, dict)
            ][:5]
        else:
            state["questions"] = []
        state["success"] = True
    except Exception as e:
        state["error"] = str(e)
        state["success"] = False
        state["should_ask_questions"] = False
        state["questions"] = []
        state["reasoning_summary"] = ""
        traceback.print_exc()
    return state


plan_workflow = _build_plan_workflow()


def plan_form_request(
    description: str,
    current_fields: Optional[List[Dict[str, Any]]] = None,
    current_title: str = "",
) -> Dict[str, Any]:
    """Run plan workflow and return should_ask_questions, questions, reasoning_summary."""
    initial: PlanState = {
        "description": description.strip(),
        "current_fields": current_fields,
        "current_title": current_title or "",
        "should_ask_questions": False,
        "questions": [],
        "reasoning_summary": "",
        "error": None,
        "success": False,
    }
    try:
        final = plan_workflow.invoke(initial)
        return {
            "should_ask_questions": final.get("should_ask_questions", False),
            "questions": final.get("questions", []),
            "reasoning_summary": final.get("reasoning_summary", "") or "Ready.",
            "success": final.get("success", False),
            "error": final.get("error"),
        }
    except Exception as e:
        return {
            "should_ask_questions": False,
            "questions": [],
            "reasoning_summary": "",
            "success": False,
            "error": str(e),
        }


workflow_instance = FormGeneratorWorkflow()