agents/graph_builder_langgraph.py

import os
import uuid
from typing import List, Dict, Optional
from pydantic import BaseModel
from langgraph.graph import StateGraph
from .document_parser import parse_documents
from .requirements_extractor import extract_requirements
from .ui_generator import generate_ui_html


# ==========================================================
# 🧱 Define the State Schema (Pydantic model)
# ==========================================================
class PipelineState(BaseModel):
    files: Optional[List[str]] = None
    text: Optional[str] = None
    requirements: Optional[Dict] = None
    html: Optional[str] = None
    public_url: Optional[str] = None


# ==========================================================
# 📄 Node 1 — Document Parser
# ==========================================================
async def node_parse_docs(state: PipelineState):
    print("📄 [Document Parser] Extracting text...")

    result = await parse_documents(state.files)

    # handle both string or dict outputs
    if isinstance(result, dict):
        text = result.get("text", "").strip()
    else:
        text = str(result).strip()

    if not text:
        raise ValueError("No text extracted from documents.")

    return {"text": text}


# ==========================================================
# 🧠 Node 2 — Requirements Extractor
# ==========================================================
async def node_requirements(state: PipelineState, config=None) -> Dict:
    print("🧩 [Requirements Agent] Extracting structured info...")

    if not state.text:
        raise ValueError("Missing 'text' field in state for requirement extraction.")

    requirements = await extract_requirements(state.text)
    if not requirements:
        raise ValueError("Requirements extraction returned empty.")

    return {"requirements": requirements}


# ==========================================================
# 🎨 Node 3 — UI Generator
# ==========================================================
async def node_ui_generator(state: PipelineState, config=None) -> Dict:
    print("🎨 [UI Generator] Creating PowerApps-style interface...")

    if not state.requirements:
        raise ValueError("Missing 'requirements' in state for UI generation.")

    # Load reference design
    ref_path = os.path.join("templates", "demo_qms_design.html")
    if not os.path.exists(ref_path):
        raise FileNotFoundError(f"Reference HTML not found at {ref_path}")

    with open(ref_path, "r", encoding="utf-8") as f:
        reference_html = f.read()

    ui = await generate_ui_html(state.requirements, reference_html)
    html = ui.get("html", "")

    if not html or "<html" not in html:
        raise ValueError("Generated HTML invalid or empty.")

    # Save to static/outputs/
    os.makedirs("static/outputs", exist_ok=True)
    session_id = uuid.uuid4().hex
    html_path = os.path.join("static/outputs", f"{session_id}.html")
    with open(html_path, "w", encoding="utf-8") as f:
        f.write(html)

    base_url = os.environ.get("BASE_PUBLIC_URL", "").rstrip("/")
    if not base_url:
        base_url = "http://localhost:7860"

    public_url = f"{base_url}/static/outputs/{session_id}.html"

    print(f"✅ [UI Generator] Generated UI available at: {public_url}")

    return {"html": html, "public_url": public_url}


# ==========================================================
# 🔗 Build the LangGraph
# ==========================================================
graph = StateGraph(PipelineState)

graph.add_node("parse_docs", node_parse_docs)
graph.add_node("requirements", node_requirements)
graph.add_node("ui_generator", node_ui_generator)

graph.set_entry_point("parse_docs")
graph.add_edge("parse_docs", "requirements")
graph.add_edge("requirements", "ui_generator")
graph.set_finish_point("ui_generator")

graph = graph.compile()

print("✅ LangGraph pipeline loaded successfully: Document → Requirements → UI Generator")