src/agents/auditor/agent.ts

import fs from "node:fs";

import { HumanMessage, SystemMessage } from "@langchain/core/messages";
import { END, type GraphNode, START, StateGraph } from "@langchain/langgraph";
import { z } from "zod";

import { createLLM } from "../../config/llm.ts";
import { emitStep, logger } from "../../logger.ts";
import { MAX_DOC_CHARS, MAX_REFLECTIONS, MAX_SOL_CHARS, MIN_FILE_IMPORTANCE } from "./config.ts";
import {
  FIND_VULNERABILITIES_PROMPT,
  GATHER_CONTEXT_PROMPT,
  JUDGE_FINDINGS_PROMPT,
  RANK_FILES_PROMPT,
  REFINE_VULNERABILITIES_PROMPT,
} from "./prompts.ts";
import { AuditorState, CandidateFindingSchema, FileRankingSchema, JudgeReviewSchema } from "./state.ts";
import { buildRepoTree } from "./tools/repo-tree/tool.ts";
import { analyzeSolidityFile } from "./tools/solidity-analyzer/tool.ts";
import { buildReviewBlocks, matchLines, walkDirectory } from "./utils.ts";

const llmHaiku = createLLM("anthropic", { model: "claude-haiku-4-5", maxTokens: 20000 });
const llmOpus = createLLM("anthropic", { model: "claude-opus-4-8", temperature: null, maxTokens: 20000 });
const llmSonnet = createLLM("anthropic", { model: "claude-sonnet-4-6", maxTokens: 20000 });

const defineScope: GraphNode<typeof AuditorState> = async (state) => {
  emitStep({ agent: "auditor", step: "scope", status: "running" });
  logger.info(`[Auditor] defineScope: percorrendo repositório em ${state.repoPath}`);

  const solFiles: string[] = [];
  const docFiles: string[] = [];

  walkDirectory(state.repoPath, 0, solFiles, docFiles);

  const fileTree = buildRepoTree(state.repoPath);

  logger.info(
    `[Auditor] defineScope: encontrado(s) ${solFiles.length} arquivo(s) Solidity e ${docFiles.length} arquivo(s) de documentação`,
  );
  logger.debug(`[Auditor] defineScope: arquivos Solidity: ${JSON.stringify(solFiles)}`);
  logger.debug(`[Auditor] defineScope: arquivos de documentação: ${JSON.stringify(docFiles)}`);
  logger.debug(`[Auditor] defineScope: árvore de arquivos:\n${fileTree}`);

  emitStep({ agent: "auditor", step: "scope", status: "done" });

  logger.info("[Auditor] defineScope: rankeando arquivos por importância");

  const RankFilesSchema = z.object({ rankings: z.array(FileRankingSchema) });
  const rankingModel = llmHaiku.withStructuredOutput(RankFilesSchema);

  const { rankings } = await rankingModel.invoke([
    new SystemMessage({ content: [{ type: "text", text: RANK_FILES_PROMPT, cache_control: { type: "ephemeral" } }] }),
    new HumanMessage(
      `Árvore de arquivos:\n\`\`\`\n${fileTree}\n\`\`\`\n\nArquivos Solidity para classificar:\n${solFiles.map((f) => `- ${f}`).join("\n")}`,
    ),
  ]);

  const sorted = [...rankings].sort((a, b) => b.importance - a.importance);
  logger.info(
    `[Auditor] defineScope: rankings:\n${sorted.map((r) => `  [${r.importance}/5] ${r.filePath} — ${r.reasoning}`).join("\n")}`,
  );

  const importantFiles = sorted.filter((r) => r.importance >= MIN_FILE_IMPORTANCE).map((r) => r.filePath);
  const skipped = solFiles.length - importantFiles.length;
  if (skipped > 0) {
    logger.info(
      `[Auditor] defineScope: pulando ${skipped} arquivo(s) de baixa importância (importância < ${MIN_FILE_IMPORTANCE})`,
    );
  }

  return { scope: importantFiles, docs: docFiles, fileTree, fileRankings: sorted };
};

const gatherContext: GraphNode<typeof AuditorState> = async (state) => {
  emitStep({ agent: "auditor", step: "ctx", status: "running" });
  logger.info(
    `[Auditor] gatherContext: processando ${state.scope.length} arquivo(s) Solidity e ${state.docs.length} arquivo(s) de documentação`,
  );

  const readFile = (filePath: string): string => {
    try {
      return fs.readFileSync(filePath, "utf-8");
    } catch {
      return "";
    }
  };

  const solidityEntries: { filePath: string; source: string; analysis: string }[] = [];
  for (const filePath of state.scope) {
    const source = readFile(filePath).slice(0, MAX_SOL_CHARS);
    if (!source) continue;

    const ranking = state.fileRankings.find((r) => r.filePath === filePath);
    const mode = ranking && ranking.importance >= 4 ? "full" : "short";
    const analysis = await analyzeSolidityFile(source, mode, filePath, ranking?.importance);
    solidityEntries.push({ filePath, source, analysis });
  }

  const docEntries: { filePath: string; content: string }[] = [];
  for (const filePath of state.docs) {
    const content = readFile(filePath).slice(0, MAX_DOC_CHARS);
    if (content) docEntries.push({ filePath, content });
  }

  const parts: string[] = [];

  if (docEntries.length > 0) {
    parts.push("## Documentação\n");
    for (const { filePath, content } of docEntries) {
      parts.push(`### ${filePath}\n${content}`);
    }
  }

  parts.push(`## Árvore de Arquivos\n\n\`\`\`\n${state.fileTree}\n\`\`\``);

  parts.push("## Análise Estrutural\n");
  for (const { analysis } of solidityEntries) {
    parts.push(analysis);
  }

  const model = llmHaiku.withStructuredOutput(z.object({ context: z.string() }));
  const result = await model.invoke([
    new SystemMessage({
      content: [{ type: "text", text: GATHER_CONTEXT_PROMPT, cache_control: { type: "ephemeral" } }],
    }),
    new HumanMessage(parts.join("\n\n")),
  ]);

  const fileTreeBlock = `## Árvore de Arquivos\n\n\`\`\`\n${state.fileTree}\n\`\`\``;
  const structuralBlock = `## Análise Estrutural dos Contratos\n\n${solidityEntries.map(({ analysis }) => analysis).join("\n\n---\n\n")}`;
  const repoContext = [result.context, fileTreeBlock, structuralBlock].join("\n\n");

  logger.debug(`[Auditor] gatherContext: contexto completo:\n${parts.join("\n\n")}`);
  logger.info(`[Auditor] gatherContext: contexto construído (${repoContext.length} caracteres)`);
  logger.debug(`[Auditor] gatherContext: contexto compactado:\n${repoContext}`);

  emitStep({ agent: "auditor", step: "ctx", status: "done" });

  return { repoContext };
};

const findVulnerabilities: GraphNode<typeof AuditorState> = async (state) => {
  const model = llmOpus.withStructuredOutput(z.object({ findings: z.array(CandidateFindingSchema) }));

  const isReflection = state.judgeReviews.length > 0;

  logger.info(
    `[Auditor] findVulnerabilities: invocando LLM para ${state.scope.length} arquivo(s) em paralelo (iteração ${state.reflectionCount + 1})`,
  );

  emitStep({ agent: "auditor", step: "find", status: "running", detail: `iter ${state.reflectionCount + 1}` });

  const cachedContext = {
    type: "text" as const,
    text: `Contexto do Protocolo:\n${state.repoContext}`,
    cache_control: { type: "ephemeral" as const },
  };

  const processFile = async (filePath: string) => {
    let source: string;
    try {
      source = fs.readFileSync(filePath, "utf-8").slice(0, MAX_SOL_CHARS);
    } catch {
      return [];
    }
    if (!source) return [];

    const fileEntries = isReflection
      ? state.candidateFindings
          .map((f, i) => ({ finding: f, review: state.judgeReviews[i] }))
          .filter(({ finding }) => finding.path === filePath)
      : [];

    const isRefinement = fileEntries.length > 0;
    const promptText = isRefinement ? REFINE_VULNERABILITIES_PROMPT : FIND_VULNERABILITIES_PROMPT;
    const contractText = isRefinement
      ? `Contrato (${filePath}):\n\n${source}\n\n${buildReviewBlocks(fileEntries, state.reflectionCount)}`
      : `Contrato (${filePath}):\n\n${source}`;

    logger.debug(`[Auditor] findVulnerabilities: processando ${filePath}`);

    const result = await model.invoke([
      new SystemMessage({ content: [{ type: "text", text: promptText, cache_control: { type: "ephemeral" } }] }),
      new HumanMessage({ content: [cachedContext, { type: "text", text: contractText }] }),
    ]);

    return result.findings.map((finding: any) => ({
      ...finding,
      path: filePath,
      location: matchLines(source, finding.codeSnippet) ?? "",
    }));
  };

  const [firstFile, ...restFiles] = state.scope;
  const firstFindings = firstFile ? await processFile(firstFile) : [];
  const restFindings = await Promise.all(restFiles.map(processFile));
  const candidateFindings = [firstFindings, ...restFindings].flat();

  logger.info(
    `[Auditor] findVulnerabilities: LLM retornou ${candidateFindings.length} finding(s) candidato(s) no total`,
  );
  logger.debug(`[Auditor] findVulnerabilities: findings:\n${JSON.stringify(candidateFindings, null, 2)}`);

  emitStep({ agent: "auditor", step: "find", status: "done" });
  return { candidateFindings };
};

const judgeFindings: GraphNode<typeof AuditorState> = async (state) => {
  emitStep({ agent: "auditor", step: "judge", status: "running" });
  if (state.candidateFindings.length === 0) {
    logger.info("[Auditor] judgeFindings: sem findings candidatos para revisar, pulando chamada ao LLM");
    emitStep({ agent: "auditor", step: "judge", status: "done" });
    return {
      judgeReviews: [],
      findings: [],
      reflectionCount: state.reflectionCount + 1,
    };
  }

  const model = llmSonnet.withStructuredOutput(JudgeReviewSchema);

  logger.info(
    `[Auditor] judgeFindings: revisando ${state.candidateFindings.length} finding(s) candidato(s) em paralelo`,
  );

  const cachedContext = {
    type: "text" as const,
    text: `Contexto do Protocolo:\n${state.repoContext}`,
    cache_control: { type: "ephemeral" as const },
  };

  const reviewFinding = async (finding: (typeof state.candidateFindings)[number], i: number) => {
    let source: string;
    try {
      source = fs.readFileSync(finding.path, "utf-8").slice(0, MAX_SOL_CHARS);
    } catch {
      source = "";
    }

    const findingText = `[Achado ${i + 1}] ${finding.title}\nSeveridade: ${finding.severity}\nDescrição: ${finding.description}\nLocalização: ${finding.path} linhas ${finding.location}\nCódigo:\n\`\`\`solidity\n${finding.codeSnippet}\n\`\`\``;

    logger.debug(`[Auditor] judgeFindings: revisando finding ${i + 1}: ${finding.title}`);
    return model.invoke([
      new SystemMessage({
        content: [{ type: "text", text: JUDGE_FINDINGS_PROMPT, cache_control: { type: "ephemeral" } }],
      }),
      new HumanMessage({
        content: [
          cachedContext,
          { type: "text", text: `Contrato (${finding.path}):\n\n${source}\n\nAchado para Revisão:\n\n${findingText}` },
        ],
      }),
    ]);
  };

  const [firstFinding, ...restFindings] = state.candidateFindings;
  const firstReview = await reviewFinding(firstFinding, 0);
  const restReviews = await Promise.all(restFindings.map((f, i) => reviewFinding(f, i + 1)));
  const reviews = [firstReview, ...restReviews];

  const confirmedEntries = state.candidateFindings
    .map((finding, i) => ({ finding, review: reviews[i] }))
    .filter(({ review }) => !review.isFalsePositive);

  const findings = confirmedEntries.map(({ finding, review }) => ({
    ...finding,
    judgeReview: {
      review: review.review,
      confidence: review.confidence,
      exploitablePaths: review.exploitablePaths,
    },
  }));

  const falsePositiveCount = state.candidateFindings.length - findings.length;

  logger.info(`[Auditor] judgeFindings: ${findings.length} confirmado(s), ${falsePositiveCount} falso(s) positivo(s)`);
  logger.debug(`[Auditor] judgeFindings: revisões:\n${JSON.stringify(reviews, null, 2)}`);

  emitStep({ agent: "auditor", step: "judge", status: "done" });
  return {
    judgeReviews: reviews,
    findings,
    reflectionCount: state.reflectionCount + 1,
  };
};

export const auditorAgent = new StateGraph(AuditorState)
  .addNode("defineScope", defineScope)
  .addNode("gatherContext", gatherContext)
  .addNode("findVulnerabilities", findVulnerabilities)
  .addNode("judgeFindings", judgeFindings)
  .addEdge(START, "defineScope")
  .addEdge("defineScope", "gatherContext")
  .addEdge("gatherContext", "findVulnerabilities")
  .addEdge("findVulnerabilities", "judgeFindings")
  .addConditionalEdges("judgeFindings", (state) => {
    const hasFalsePositives = state.judgeReviews.some((r) => r.isFalsePositive);
    if (hasFalsePositives && state.reflectionCount < MAX_REFLECTIONS) {
      return "findVulnerabilities";
    }
    return END;
  })
  .compile();

export const testAgent = new StateGraph(AuditorState)
  .addNode("defineScope", defineScope)
  .addNode("gatherContext", gatherContext)
  .addEdge(START, "defineScope")
  .addEdge("defineScope", "gatherContext")
  .addEdge("gatherContext", END)
  .compile();