// ============================================================================
// LangGraph Node Definitions
// Each node wraps an agent and handles state transitions
// ============================================================================

import { AgentState, AgentName } from '../types';
import { scannerAgent } from '../agents/scanner.agent';
import { analyzerAgent } from '../agents/analyzer.agent';
import { strategistAgent } from '../agents/strategist.agent';
import { writerAgent } from '../agents/writer.agent';
import { reviewerAgent } from '../agents/reviewer.agent';
import { runnerAgent } from '../agents/runner.agent';
import { securityAgent } from '../agents/security.agent';
import { reporterAgent } from '../agents/reporter.agent';
import { mergeState } from './state';

export type NodeFunction = (state: AgentState) => Promise<AgentState>;

/**
 * Creates a graph node that wraps an agent function.
 * Handles state merging and error propagation.
 */
function createNode(agent: AgentName, fn: (state: AgentState) => Promise<Partial<AgentState>>): NodeFunction {
  return async (state: AgentState): Promise<AgentState> => {
    const updatedState = mergeState(state, { currentAgent: agent, status: 'running' });

    try {
      const result = await fn(updatedState);
      return mergeState(updatedState, result);
    } catch (error) {
      const errMsg = error instanceof Error ? error.message : String(error);
      return mergeState(updatedState, {
        errors: [...updatedState.errors, `${agent}: ${errMsg}`],
        status: 'failed',
      });
    }
  };
}

// --- Graph Nodes ---

export const scannerNode = createNode('scanner', scannerAgent);
export const analyzerNode = createNode('analyzer', analyzerAgent);
export const strategistNode = createNode('strategist', strategistAgent);
export const writerNode = createNode('writer', writerAgent);
export const reviewerNode = createNode('reviewer', reviewerAgent);
export const runnerNode = createNode('runner', runnerAgent);
export const securityNode = createNode('security', securityAgent);
export const reporterNode = createNode('reporter', reporterAgent);

// --- Conditional Edges ---

/**
 * Determines if the workflow should continue after the reviewer.
 * If review score is too low, go back to writer for improvements.
 */
export function shouldRewrite(state: AgentState): 'runner' | 'writer' {
  if (state.testReviews.length === 0) return 'runner';

  const avgScore = state.testReviews.reduce((s, r) => s + r.score, 0) / state.testReviews.length;
  const criticalErrors = state.testReviews.reduce(
    (s, r) => s + r.issues.filter(i => i.severity === 'error').length, 0
  );

  // If too many critical errors, rewrite (but only once to avoid infinite loop)
  const rewriteAttempts = state.agentLog.filter(l => l.agent === 'writer').length;
  if (criticalErrors > 3 && rewriteAttempts < 2) {
    return 'writer';
  }

  return 'runner';
}

/**
 * Determines if the workflow should continue or stop after errors.
 */
export function shouldContinue(state: AgentState): 'continue' | 'stop' {
  if (state.status === 'failed') return 'stop';
  if (state.errors.length > 10) return 'stop';
  return 'continue';
}

/**
 * All available nodes mapped by name.
 */
export const nodeMap: Record<AgentName, NodeFunction> = {
  scanner: scannerNode,
  analyzer: analyzerNode,
  strategist: strategistNode,
  writer: writerNode,
  reviewer: reviewerNode,
  runner: runnerNode,
  security: securityNode,
  reporter: reporterNode,
};