import { NextResponse } from "next/server";
import { ulid } from "ulid";
import { Command } from "@langchain/langgraph";
import { isAIMessage } from "@langchain/core/messages";
import checkpointer from "@/server/checkpointer";
import { PrivyClient } from '@privy-io/server-auth';
import { PublicKey } from '@solana/web3.js';
import { SolanaAgentKit } from '@/solana-agent-kit';
import pool from "@/server/db";
import { createPrivyEmbeddedWallet } from '@/lib/solana/PrivyEmbeddedWallet';
import getGraph from "@/server/graph";

export const dynamic = "force-dynamic";

const PRIVY_APP_ID = process.env.NEXT_PUBLIC_PRIVY_APP_ID;
const PRIVY_APP_SECRET = process.env.PRIVY_APP_SECRET;
const PRIVY_WALLET_SIGNING_KEY = process.env.PRIVY_WALLET_SIGNING_KEY;
const RPC_URL = process.env.RPC_URL;
const OPENAI_API_KEY = process.env.OPENAI_API_KEY;
const HELIUS_API_KEY = process.env.HELIUS_API_KEY;

// TODO
const SYSTEMprompt = `You are a senior Solana software engineer and market data expert. You have access to Solana Defi information. You could also analyze github repos and tweet profiles and posts and search twitter. 
Your goal is to give you clean and concise responses to users to help them on your Solana journey. You must analyze carefully the user input, then reason about the possible routes you could take using the tools available to you, dont return information not asked for by the user. All tool calls are visible to the user so DO NOT repeat what is said in the tools.
YOU MUST draw your own conclusions and analysis from the tool responses. Only your analysis must be show the user. Keep it concise and small. If user asks for list of tokens trending the charts will be showm, DO NOT LIST RETURN THEM TO THE USER! KEEP RESPONSES SHORT!`;

export async function POST(req) {
  try {
    const body = await req.json();
    console.log(body);

    const { message, thread_id, approveObj = {} } = body;
    console.log(thread_id);

    if (!thread_id) {
      return NextResponse.json({ error: "Unauthorized. Please sign in" }, { status: 401 });
    }
    if (!message && approveObj == {}) {
      return NextResponse.json({ error: "Empty message" }, { status: 401 });
    }


    const cookieAuthToken = req.cookies.get("privy-id-token");
    console.log(cookieAuthToken);
    
    if (!cookieAuthToken) {
      console.log("No authentication token found.");
      return NextResponse.json({ error: "Unauthorized: No auth token" }, { status: 401 });
    }
  
      const PRIVY_SERVER_CLIENT = new PrivyClient(PRIVY_APP_ID, PRIVY_APP_SECRET, {
        walletApi: {
          authorizationPrivateKey: PRIVY_WALLET_SIGNING_KEY,
        },
      });
      
      console.log("Verifying auth token...");
      const claims = await PRIVY_SERVER_CLIENT.verifyAuthToken(cookieAuthToken.value);
      const userId = claims.userId;
      console.log("Authenticated user ID:", userId);
  
      const dbRes = await pool.query("SELECT userid, publickey, tg_userid, delegated FROM users WHERE userid = $1", [userId]);
      if (dbRes.rows.length === 1) {
        const { userid, publickey } = dbRes.rows[0];
        console.log("User found in database:", userid, "with public key:", publickey);
  
        const walletAdapter = new createPrivyEmbeddedWallet(
          PRIVY_SERVER_CLIENT,
          new PublicKey(publickey)
        );
  
        const solAgentKit = new SolanaAgentKit(walletAdapter, RPC_URL, {
          OPENAI_API_KEY,
          HELIUS_API_KEY
        });
        const graph = getGraph(solAgentKit)

        const input = {
          role: "user",
          content: message ?? "",
        };
    
        const config = { configurable: { thread_id } };
        let passToGraph = { messages: [input] };
    
        const checkpointerData = await checkpointer.get(config);
    
        if (
          typeof checkpointerData?.channel_values === "object" &&
          "branch:agent:condition:checkApproval" in checkpointerData.channel_values
        ) {
          console.log("=".repeat(75));
          console.log(`${"=".repeat(25)} Graph interrupted for user input ${"=".repeat(25)}`);
          if (Object.keys(approveObj).length == 0) {
            return NextResponse.json({ error: "Please accept or reject the tools" }, { status: 401 });
          }
          passToGraph = new Command({ resume: approveObj });
        }
    
        const transformStream = new ReadableStream({
          async start(controller) {
            const textEncoder = new TextEncoder();
            try {
              for await (const event of await graph.stream(passToGraph, config)) {
                if (event.__interrupt__) {
                  console.log("=".repeat(75));
                  console.log("=" + "=".repeat(25) + " Interrupted for permission " + "=".repeat(25));
                  console.log(JSON.stringify(event.__interrupt__));
                  console.log("=".repeat(75));
                  const toolCalls = event.__interrupt__[0].value.map((tool) => ({
                    id: tool.id,
                    name: tool.name,
                    args: Object.fromEntries(Object.entries(tool.args).filter(([key]) => key !== "debug")),
                  }));
                  console.log(toolCalls);
    
                  controller.enqueue(`f:{"messageId":"${event.__interrupt__[0].ns[0]}"}\n`);
                  const interrupter_tool_id = ulid();
                  controller.enqueue(
                    `9:${JSON.stringify({
                      toolCallId: interrupter_tool_id,
                      toolName: "interrupter",
                      args: { toolCalls },
                    })}\n`,
                  );
    
                  controller.enqueue(
                    `a:${JSON.stringify({
                      toolCallId: interrupter_tool_id,
                      result: "Your input is required to proceed",
                    })}\n`,
                  );
    
                  controller.enqueue(
                    `e:${JSON.stringify({
                      finishReason: "tool-calls",
                      usage: { promptTokens: 0, completionTokens: 0 },
                      isContinued: false,
                    })}\n`,
                  );
                } else if (event.checkApproval) {
                  continue;
                } else if (event.agent) {
                  console.log("=".repeat(75));
                  console.log("=" + "=".repeat(25) + " Agent node " + "=".repeat(25));
                  const message = event.agent.messages[0];
                  console.log(message);
    
                  if (message.content) {
                    controller.enqueue(`f:{"messageId":"${message.id}"}\n`);
                    controller.enqueue(`0:${JSON.stringify(message.content)}\n`);
                    controller.enqueue(
                      `e:${JSON.stringify({
                        finishReason: "stop",
                        usage: {
                          promptTokens: message.usage_metadata.input_tokens,
                          completionTokens: message.usage_metadata.output_tokens,
                        },
                        isContinued: false,
                      })}\n`,
                    );
                  }
                } else if (event.allToolsNode) {
                  console.log("=".repeat(75));
                  console.log("=" + "=".repeat(25) + " All tools node " + "=".repeat(25));
                  console.log(event.allToolsNode);
                  console.log("=".repeat(75));
                  for (const toolMessage of event.allToolsNode.messages) {
                    controller.enqueue(`f:{"messageId":"${toolMessage.id}"}\n`);
    
                    controller.enqueue(
                      textEncoder.encode(
                        `9:${JSON.stringify({
                          toolCallId: toolMessage.tool_call_id,
                          toolName: toolMessage.name,
                          args: { artifact: toolMessage.artifact },
                        })}\n`,
                      ),
                    );
    
                    controller.enqueue(
                      textEncoder.encode(
                        `a:${JSON.stringify({
                          toolCallId: toolMessage.tool_call_id,
                          result: toolMessage.content,
                        })}\n`,
                      ),
                    );
    
                    controller.enqueue(
                      `e:${JSON.stringify({
                        finishReason: "tool-calls",
                        usage: { promptTokens: 0, completionTokens: 0 },
                        isContinued: false,
                      })}\n`,
                    );
                  }
                }
              }
              const checkpointerData = await checkpointer.get(config);
    
              if (
                typeof checkpointerData?.channel_values === "object" &&
                !("branch:agent:condition:checkApproval" in checkpointerData.channel_values)
              ) {
                const lastMessage = checkpointerData?.channel_values?.messages?.at(-1);
                const lastMessage_hasToolCalls = lastMessage?.tool_calls?.length > 0;
                const lastMessage_isAIMessageFlag = lastMessage ? isAIMessage(lastMessage) : false;
    
                if (lastMessage_isAIMessageFlag && !lastMessage_hasToolCalls) {
                  console.log("Breaking loop: Last message is an AI message with no tool calls.");
                  // console.log(checkpointerData);
                  controller.enqueue(`d:{"finishReason":"stop","usage":{"promptTokens":0,"completionTokens":0}}\n`);
                }
              }
            } catch (error) {
              controller.enqueue(`3:${JSON.stringify(error.message)}\n`);
              controller.error(error);
            } finally {
              controller.close();
            }
          },
        });
    
        const response = new Response(transformStream, {
          headers: {
            "Content-Type": "text/event-stream", 
            "Cache-Control": "no-cache", 
            Connection: "keep-alive", 
            "x-vercel-ai-data-stream": "v1",
          },
        });
    
        return response;
    


      } else {
        console.log("User not found in database.");
        return NextResponse.json({ error: "User not found" }, { status: 404 });
      }
  } catch (e) {
    console.log(e);
    return NextResponse.json({ error: "Error while processing your request. Please try again later" }, { status: 500 });
  }
}