import { Annotation, StateGraph, START, END } from "@langchain/langgraph";
import type { AnalysisResult } from "@/lib/ai/smart-analyzer";
import type { RouteResult } from "@/lib/ai/model-router";

interface HistoryMessage {
  role: "user" | "assistant";
  content: string;
}

export interface AutoAgentInput {
  text: string;
  hasImage: boolean;
  imageUrls?: string[];
  imageFromHistory?: boolean;
  historyMessages?: HistoryMessage[];
  context?: string;
}

export interface AutoAgentResult {
  analysis?: AnalysisResult;
  friendlyMessage?: string;
  optimizedPrompt?: string;
  route?: RouteResult;
}

const AutoAgentState = Annotation.Root({
  text: Annotation<string>,
  hasImage: Annotation<boolean>,
  imageUrls: Annotation<string[]>,
  imageFromHistory: Annotation<boolean>,
  historyMessages: Annotation<HistoryMessage[]>,
  context: Annotation<string | null>,
  analysis: Annotation<AnalysisResult | null>,
  friendlyMessage: Annotation<string | null>,
  optimizedPrompt: Annotation<string | null>,
  route: Annotation<RouteResult | null>,
});

type AutoAgentStateType = typeof AutoAgentState.State;

const analyzeNode = async (state: AutoAgentStateType) => {
  const { analyzeIntent } = await import("./smart-analyzer");
  const analysis = await analyzeIntent({
    text: state.text,
    hasImage: state.hasImage,
    imageUrls: state.imageUrls,
    imageFromHistory: state.imageFromHistory,
    historyMessages: state.historyMessages,
    context: state.context ?? undefined,
  });

  const fallbackQuestion = analysis.needsClarification
    ? analysis.clarifyingQuestions?.[0]
    : undefined;

  return {
    analysis,
    friendlyMessage: analysis.conversationalReply || fallbackQuestion || null,
  };
};

const shouldOptimize = (state: AutoAgentStateType) => {
  if (!state.analysis) return "route";
  if (state.analysis.needsClarification) return "route";
  if (state.analysis.intent === "image_edit" && state.analysis.subIntent) {
    return "optimize";
  }
  return "route";
};

const optimizeNode = async (state: AutoAgentStateType) => {
  if (!state.analysis || state.analysis.intent !== "image_edit" || !state.analysis.subIntent) {
    return {};
  }
  const { optimizePromptWithLLM } = await import("./prompt-optimizer");
  const optimized = await optimizePromptWithLLM({
    originalPrompt: state.text,
    intent: state.analysis.intent,
    subIntent: state.analysis.subIntent,
    historyMessages: state.historyMessages,
    context: state.context ?? undefined,
  });
  return { optimizedPrompt: optimized.optimizedPrompt };
};

const routeNode = async (state: AutoAgentStateType) => {
  const { routeToModel } = await import("./model-router");
  const route = await routeToModel({
    intent: state.analysis?.intent || "chat",
    subIntent: state.analysis?.subIntent,
    hasImage: state.hasImage,
  });
  return { route };
};

let compiledGraph: ReturnType<typeof buildGraph> | null = null;

function buildGraph() {
  return new StateGraph(AutoAgentState)
    .addNode("analyze", analyzeNode)
    .addNode("optimize", optimizeNode)
    .addNode("route", routeNode)
    .addEdge(START, "analyze")
    .addConditionalEdges("analyze", shouldOptimize, {
      optimize: "optimize",
      route: "route",
    })
    .addEdge("optimize", "route")
    .addEdge("route", END)
    .compile();
}

function getGraph() {
  if (!compiledGraph) {
    compiledGraph = buildGraph();
  }
  return compiledGraph;
}

export async function runAutoAgentGraph(params: AutoAgentInput): Promise<AutoAgentResult> {
  const graph = getGraph();
  const result = await graph.invoke({
    text: params.text,
    hasImage: params.hasImage,
    imageUrls: params.imageUrls ?? [],
    imageFromHistory: params.imageFromHistory ?? false,
    historyMessages: params.historyMessages ?? [],
    context: params.context ?? null,
    analysis: null,
    friendlyMessage: null,
    optimizedPrompt: null,
    route: null,
  });

  return {
    analysis: result.analysis ?? undefined,
    friendlyMessage: result.friendlyMessage ?? undefined,
    optimizedPrompt: result.optimizedPrompt ?? undefined,
    route: result.route ?? undefined,
  };
}