app/api/routes/graph.py

"""
Graph API Routes.

Endpoints for creating, managing, and executing workflow graphs.
"""

from typing import Any, Dict, Optional
from fastapi import APIRouter, HTTPException, BackgroundTasks, status
from uuid import uuid4
import logging

from app.api.schemas import (
    GraphCreateRequest,
    GraphCreateResponse,
    GraphRunRequest,
    GraphRunResponse,
    GraphInfoResponse,
    GraphListResponse,
    RunStateResponse,
    RunListResponse,
    ExecutionLogEntry,
    ExecutionStatus,
    ErrorResponse,
)
from app.engine.graph import Graph, END
from app.engine.node import Node, get_registered_node
from app.engine.executor import Executor, ExecutionResult
from app.storage.memory import graph_storage, run_storage
from app.tools.registry import tool_registry


logger = logging.getLogger(__name__)

router = APIRouter(prefix="/graph", tags=["Graph"])


# ============================================================
# Condition Functions Registry
# ============================================================

# Built-in condition functions for routing
_condition_registry: Dict[str, Any] = {}


def register_condition(name: str):
    """Decorator to register a condition function."""
    def decorator(func):
        _condition_registry[name] = func
        return func
    return decorator


@register_condition("quality_check")
def quality_check_condition(state: Dict[str, Any]) -> str:
    """Route based on quality score vs threshold."""
    quality_score = state.get("quality_score", 0)
    threshold = state.get("quality_threshold", 7.0)
    return "pass" if quality_score >= threshold else "fail"


# Also register as quality_meets_threshold (used by code review workflow)
@register_condition("quality_meets_threshold")
def quality_meets_threshold(state: Dict[str, Any]) -> str:
    """Route based on quality score vs threshold."""
    quality_score = state.get("quality_score", 0)
    threshold = state.get("quality_threshold", 7.0)
    return "pass" if quality_score >= threshold else "fail"


@register_condition("always_continue")
def always_continue(state: Dict[str, Any]) -> str:
    """Always returns 'continue' - for unconditional looping."""
    return "continue"


# Also register as always_loop (used by code review workflow)
@register_condition("always_loop")
def always_loop(state: Dict[str, Any]) -> str:
    """Always returns 'continue' - for looping back."""
    return "continue"


@register_condition("always_end")
def always_end(state: Dict[str, Any]) -> str:
    """Always returns 'end' - for explicit termination."""
    return "end"


@register_condition("max_iterations_check")
def max_iterations_check(state: Dict[str, Any]) -> str:
    """Check if max iterations reached."""
    iteration = state.get("_iteration", 0)
    max_iter = state.get("_max_iterations", 3)
    return "stop" if iteration >= max_iter else "continue"


def get_condition(name: str):
    """Get a condition function by name."""
    return _condition_registry.get(name)


# ============================================================
# Graph CRUD Endpoints
# ============================================================

@router.post(
    "/create",
    response_model=GraphCreateResponse,
    status_code=status.HTTP_201_CREATED,
    responses={
        400: {"model": ErrorResponse, "description": "Invalid graph definition"},
        404: {"model": ErrorResponse, "description": "Handler not found"},
    }
)
async def create_graph(request: GraphCreateRequest) -> GraphCreateResponse:
    """
    Create a new workflow graph.
    
    Define nodes with their handlers, edges for flow control,
    and conditional edges for branching logic.
    """
    graph_id = str(uuid4())
    
    # Build the graph
    graph = Graph(
        graph_id=graph_id,
        name=request.name,
        description=request.description or "",
        max_iterations=request.max_iterations,
    )
    
    # Add nodes
    for node_def in request.nodes:
        # Find the handler function
        handler = get_registered_node(node_def.handler)
        if handler is None:
            # Check tool registry as fallback
            tool = tool_registry.get(node_def.handler)
            if tool:
                handler = _create_node_handler_from_tool(node_def.handler)
            else:
                raise HTTPException(
                    status_code=404,
                    detail=f"Handler '{node_def.handler}' not found. "
                           f"Available handlers: {list(tool_registry.list_tools())}"
                )
        
        graph.add_node(
            name=node_def.name,
            handler=handler,
            description=node_def.description or "",
        )
    
    # Add direct edges
    for source, target in request.edges.items():
        if source not in graph.nodes:
            raise HTTPException(
                status_code=400,
                detail=f"Edge source '{source}' is not a valid node"
            )
        if target != END and target != "__END__" and target not in graph.nodes:
            raise HTTPException(
                status_code=400,
                detail=f"Edge target '{target}' is not a valid node"
            )
        # Normalize END
        target = END if target == "__END__" else target
        graph.add_edge(source, target)
    
    # Add conditional edges
    for source, cond_routes in request.conditional_edges.items():
        if source not in graph.nodes:
            raise HTTPException(
                status_code=400,
                detail=f"Conditional edge source '{source}' is not a valid node"
            )
        
        # Get condition function
        condition_func = get_condition(cond_routes.condition)
        if condition_func is None:
            raise HTTPException(
                status_code=404,
                detail=f"Condition '{cond_routes.condition}' not found. "
                       f"Available: {list(_condition_registry.keys())}"
            )
        
        # Normalize routes (handle __END__)
        routes = {}
        for key, target in cond_routes.routes.items():
            if target == "__END__":
                routes[key] = END
            else:
                if target not in graph.nodes:
                    raise HTTPException(
                        status_code=400,
                        detail=f"Conditional route target '{target}' is not a valid node"
                    )
                routes[key] = target
        
        graph.add_conditional_edge(source, condition_func, routes)
    
    # Set entry point
    if request.entry_point:
        if request.entry_point not in graph.nodes:
            raise HTTPException(
                status_code=400,
                detail=f"Entry point '{request.entry_point}' is not a valid node"
            )
        graph.set_entry_point(request.entry_point)
    
    # Validate graph
    errors = graph.validate()
    if errors:
        raise HTTPException(
            status_code=400,
            detail=f"Graph validation failed: {errors}"
        )
    
    # Store the graph
    await graph_storage.save(
        graph_id=graph_id,
        name=request.name,
        definition=graph.to_dict(),
    )
    
    logger.info(f"Created graph: {graph_id} ({request.name})")
    
    return GraphCreateResponse(
        graph_id=graph_id,
        name=request.name,
        message="Graph created successfully",
        node_count=len(graph.nodes),
    )


def _create_node_handler_from_tool(tool_name: str):
    """Create a node handler that calls a tool and updates state."""
    def handler(state: Dict[str, Any]) -> Dict[str, Any]:
        tool = tool_registry.get(tool_name)
        if not tool:
            raise ValueError(f"Tool '{tool_name}' not found")
        
        # Check if the tool function expects a 'state' parameter (node handler style)
        # or individual parameters (regular tool style)
        import inspect
        sig = inspect.signature(tool.func)
        param_names = list(sig.parameters.keys())
        
        if len(param_names) == 1 and param_names[0] == 'state':
            # This is a node handler - pass state directly
            result = tool.func(state)
        else:
            # This is a regular tool - extract arguments from state
            result = tool.func(**_extract_tool_args(tool, state))
        
        # Handle the result
        if isinstance(result, dict):
            # If the tool returns a full state, use it directly
            # Check if it looks like a state update (has same keys or adds new ones)
            if result is state:
                return result
            # Merge result into state
            state.update(result)
        
        return state
    
    handler.__name__ = f"{tool_name}_handler"
    return handler


def _extract_tool_args(tool, state: Dict[str, Any]) -> Dict[str, Any]:
    """Extract arguments for a tool from state."""
    import inspect
    sig = inspect.signature(tool.func)
    args = {}
    
    for param_name, param in sig.parameters.items():
        if param_name in state:
            args[param_name] = state[param_name]
        elif param.default != inspect.Parameter.empty:
            pass  # Use default
        # Skip missing optional params
    
    return args


@router.get(
    "/{graph_id}",
    response_model=GraphInfoResponse,
    responses={404: {"model": ErrorResponse}},
)
async def get_graph(graph_id: str) -> GraphInfoResponse:
    """Get information about a specific graph."""
    stored = await graph_storage.get(graph_id)
    if not stored:
        raise HTTPException(status_code=404, detail=f"Graph '{graph_id}' not found")
    
    definition = stored.definition
    
    # Generate mermaid diagram
    mermaid = _generate_mermaid(definition)
    
    return GraphInfoResponse(
        graph_id=stored.graph_id,
        name=stored.name,
        description=definition.get("description"),
        node_count=len(definition.get("nodes", {})),
        nodes=list(definition.get("nodes", {}).keys()),
        entry_point=definition.get("entry_point"),
        max_iterations=definition.get("max_iterations", 100),
        created_at=stored.created_at.isoformat(),
        mermaid_diagram=mermaid,
    )


def _generate_mermaid(definition: Dict[str, Any]) -> str:
    """Generate a Mermaid diagram from graph definition."""
    lines = ["graph TD"]
    
    nodes = definition.get("nodes", {})
    edges = definition.get("edges", {})
    cond_edges = definition.get("conditional_edges", {})
    
    # Add nodes
    for name in nodes:
        label = name.replace("_", " ").title()
        lines.append(f'    {name}["{label}"]')
    
    # Check if END is used
    has_end = END in edges.values()
    for cond in cond_edges.values():
        if END in cond.get("routes", {}).values():
            has_end = True
    
    if has_end:
        lines.append(f'    {END}(("END"))')
    
    # Add direct edges
    for source, target in edges.items():
        lines.append(f"    {source} --> {target}")
    
    # Add conditional edges
    for source, cond in cond_edges.items():
        for route_key, target in cond.get("routes", {}).items():
            lines.append(f"    {source} -->|{route_key}| {target}")
    
    return "\n".join(lines)


@router.get(
    "/",
    response_model=GraphListResponse,
)
async def list_graphs() -> GraphListResponse:
    """List all available graphs."""
    graphs = await graph_storage.list_all()
    
    graph_infos = []
    for stored in graphs:
        definition = stored.definition
        graph_infos.append(GraphInfoResponse(
            graph_id=stored.graph_id,
            name=stored.name,
            description=definition.get("description"),
            node_count=len(definition.get("nodes", {})),
            nodes=list(definition.get("nodes", {}).keys()),
            entry_point=definition.get("entry_point"),
            max_iterations=definition.get("max_iterations", 100),
            created_at=stored.created_at.isoformat(),
            mermaid_diagram=None,  # Skip for list view
        ))
    
    return GraphListResponse(graphs=graph_infos, total=len(graph_infos))


@router.delete(
    "/{graph_id}",
    status_code=status.HTTP_204_NO_CONTENT,
    responses={404: {"model": ErrorResponse}},
)
async def delete_graph(graph_id: str):
    """Delete a graph."""
    deleted = await graph_storage.delete(graph_id)
    if not deleted:
        raise HTTPException(status_code=404, detail=f"Graph '{graph_id}' not found")
    logger.info(f"Deleted graph: {graph_id}")


# ============================================================
# Execution Endpoints
# ============================================================

@router.post(
    "/run",
    response_model=GraphRunResponse,
    responses={
        404: {"model": ErrorResponse},
        500: {"model": ErrorResponse, "description": "Execution failed"},
    }
)
async def run_graph(
    request: GraphRunRequest,
    background_tasks: BackgroundTasks,
) -> GraphRunResponse:
    """
    Execute a workflow graph with the given initial state.
    
    If `async_execution` is True, the workflow runs in the background
    and you can poll the status using GET /graph/state/{run_id}.
    """
    # Get the graph
    stored = await graph_storage.get(request.graph_id)
    if not stored:
        raise HTTPException(
            status_code=404,
            detail=f"Graph '{request.graph_id}' not found"
        )
    
    # Rebuild the graph from definition
    graph = await _rebuild_graph_from_definition(stored.definition)
    
    # Create run
    run_id = str(uuid4())
    await run_storage.create(run_id, request.graph_id, request.initial_state)
    
    if request.async_execution:
        # Run in background
        background_tasks.add_task(
            _execute_in_background,
            graph,
            run_id,
            request.initial_state,
        )
        
        return GraphRunResponse(
            run_id=run_id,
            graph_id=request.graph_id,
            status=ExecutionStatus.PENDING,
            final_state={},
            execution_log=[],
            started_at=None,
            completed_at=None,
            total_duration_ms=None,
            iterations=0,
        )
    
    # Execute synchronously
    try:
        executor = Executor(
            graph,
            run_id=run_id,
            on_step=lambda step, state: _update_run_state(run_id, step, state),
        )
        result = await executor.run(request.initial_state)
        
        # Update storage
        if result.status.value == "completed":
            await run_storage.complete(
                run_id,
                result.final_state,
                [s.to_dict() for s in result.execution_log],
            )
        else:
            await run_storage.fail(run_id, result.error or "Unknown error", result.final_state)
        
        return _result_to_response(result)
        
    except Exception as e:
        logger.exception(f"Execution failed: {e}")
        await run_storage.fail(run_id, str(e))
        raise HTTPException(status_code=500, detail=str(e))


async def _rebuild_graph_from_definition(definition: Dict[str, Any]) -> Graph:
    """Rebuild a Graph object from its stored definition."""
    graph = Graph(
        graph_id=definition.get("graph_id", str(uuid4())),
        name=definition.get("name", "Unnamed"),
        description=definition.get("description", ""),
        max_iterations=definition.get("max_iterations", 100),
    )
    
    # Add nodes
    nodes_def = definition.get("nodes", {})
    for node_name, node_info in nodes_def.items():
        handler_name = node_info.get("handler", node_name)
        handler = _create_node_handler_from_tool(handler_name)
        graph.add_node(
            name=node_name,
            handler=handler,
            description=node_info.get("description", ""),
        )
    
    # Add direct edges
    for source, target in definition.get("edges", {}).items():
        graph.add_edge(source, target)
    
    # Add conditional edges
    for source, cond_info in definition.get("conditional_edges", {}).items():
        condition_name = cond_info.get("condition", "always_continue")
        condition_func = get_condition(condition_name)
        if condition_func is None:
            condition_func = always_continue
        
        routes = cond_info.get("routes", {})
        graph.add_conditional_edge(source, condition_func, routes)
    
    # Set entry point
    if definition.get("entry_point"):
        graph.set_entry_point(definition["entry_point"])
    
    return graph


async def _execute_in_background(graph: Graph, run_id: str, initial_state: Dict[str, Any]):
    """Execute a workflow in the background."""
    try:
        executor = Executor(
            graph,
            run_id=run_id,
            on_step=lambda step, state: _update_run_state(run_id, step, state),
        )
        result = await executor.run(initial_state)
        
        if result.status.value == "completed":
            await run_storage.complete(
                run_id,
                result.final_state,
                [s.to_dict() for s in result.execution_log],
            )
        else:
            await run_storage.fail(run_id, result.error or "Unknown error", result.final_state)
            
    except Exception as e:
        logger.exception(f"Background execution failed: {e}")
        await run_storage.fail(run_id, str(e))


def _update_run_state(run_id: str, step, state: Dict[str, Any]):
    """Update run state during execution (sync callback)."""
    import asyncio
    try:
        loop = asyncio.get_event_loop()
        if loop.is_running():
            asyncio.create_task(
                run_storage.update_state(run_id, state, step.node, step.iteration)
            )
    except Exception:
        pass  # Ignore errors in callback


def _result_to_response(result: ExecutionResult) -> GraphRunResponse:
    """Convert ExecutionResult to API response."""
    return GraphRunResponse(
        run_id=result.run_id,
        graph_id=result.graph_id,
        status=ExecutionStatus(result.status.value),
        final_state=result.final_state,
        execution_log=[
            ExecutionLogEntry(
                step=s.step,
                node=s.node,
                started_at=s.started_at.isoformat(),
                completed_at=s.completed_at.isoformat() if s.completed_at else None,
                duration_ms=s.duration_ms,
                iteration=s.iteration,
                result=s.result,
                error=s.error,
                route_taken=s.route_taken,
            )
            for s in result.execution_log
        ],
        started_at=result.started_at.isoformat() if result.started_at else None,
        completed_at=result.completed_at.isoformat() if result.completed_at else None,
        total_duration_ms=result.total_duration_ms,
        iterations=result.iterations,
        error=result.error,
    )


# ============================================================
# Run State Endpoints
# ============================================================

@router.get(
    "/state/{run_id}",
    response_model=RunStateResponse,
    responses={404: {"model": ErrorResponse}},
)
async def get_run_state(run_id: str) -> RunStateResponse:
    """
    Get the current state of a workflow run.
    
    Use this to poll the status of async executions.
    """
    stored = await run_storage.get(run_id)
    if not stored:
        raise HTTPException(status_code=404, detail=f"Run '{run_id}' not found")
    
    return RunStateResponse(
        run_id=stored.run_id,
        graph_id=stored.graph_id,
        status=ExecutionStatus(stored.status),
        current_node=stored.current_node,
        current_state=stored.current_state,
        iteration=stored.iteration,
        execution_log=[
            ExecutionLogEntry(**entry) for entry in stored.execution_log
        ],
        started_at=stored.started_at.isoformat(),
        completed_at=stored.completed_at.isoformat() if stored.completed_at else None,
        error=stored.error,
    )


@router.get(
    "/runs",
    response_model=RunListResponse,
)
async def list_runs(graph_id: Optional[str] = None) -> RunListResponse:
    """List all runs, optionally filtered by graph_id."""
    if graph_id:
        runs = await run_storage.list_by_graph(graph_id)
    else:
        runs = await run_storage.list_all()
    
    run_states = []
    for stored in runs:
        run_states.append(RunStateResponse(
            run_id=stored.run_id,
            graph_id=stored.graph_id,
            status=ExecutionStatus(stored.status),
            current_node=stored.current_node,
            current_state=stored.current_state,
            iteration=stored.iteration,
            execution_log=[
                ExecutionLogEntry(**entry) for entry in stored.execution_log
            ],
            started_at=stored.started_at.isoformat(),
            completed_at=stored.completed_at.isoformat() if stored.completed_at else None,
            error=stored.error,
        ))
    
    return RunListResponse(runs=run_states, total=len(run_states))