const { ChatOpenAI } = require("@langchain/openai");
const { awaitAllCallbacks } = require("@langchain/core/callbacks/promises");
const { Calculator } = require("@langchain/community/tools/calculator");
const { AgentExecutor, createToolCallingAgent } = require("langchain/agents");
const { ChatPromptTemplate } = require("@langchain/core/prompts");
const { JLINCTracer, JLINCAuthDecision, JLINCAuthBaseChatModel, JLINCAuthTool } = require("../src/index.js");

class CalculatorPrivate extends Calculator {
  static lc_name() {
    return "CalculatorPrivate";
  }
}

class CalculatorPublic extends Calculator {
  static lc_name() {
    return "CalculatorPublic";
  }
}

async function main() {
  const config = {
    dataStoreApiUrl: "http://localhost:9090",
    dataStoreApiKey: process.env.JLINC_DATA_STORE_API_KEY,
    archiveApiUrl: "http://localhost:9090",
    archiveApiKey: process.env.JLINC_ARCHIVE_API_KEY,
    agreementId: "00000000-0000-0000-0000-000000000000",
    systemPrefix: "TestTracerJlinc",
    debug: true,
  }
 
  const jlincAuthDecision = new JLINCAuthDecision(config);
  const auth = {
    subject: {
      type: "user",
      id: "tester",
    },
    action: {
      name: "read",
    },
    resource: {
      type: "data",
      id: "1234",
      properties: {
        ownerID: "tester@test.com",
      }
    }
  }
  await jlincAuthDecision.evaluate(auth);

  const tracer = new JLINCTracer(config);

  const authorizedLlm = new ChatOpenAI({
    openAIApiKey: "n/a",
    configuration: {
      baseURL: "http://localhost:1234/v1",
    },
    modelName: "meta-llama-3.1-8b-instruct",
  });
  const notAuthorizedLlm = new ChatOpenAI({
    openAIApiKey: "n/a",
    configuration: {
      baseURL: "http://localhost:1234/v1",
    },
    modelName: "hermes-3-llama-3.1-8b",
  });
  const llm = new JLINCAuthBaseChatModel({
    config,
    jlincAuthDecision,
    targetAuthorized: authorizedLlm,
    targetNotAuthorized: notAuthorizedLlm, // Optional
  });

  const calculatorPublic = new CalculatorPublic();
  calculatorPublic.name = 'calculator_public';
  const calculatorPrivate = new CalculatorPrivate();
  calculatorPrivate.name = 'calculator_private';
  const jlincAuthTool = new JLINCAuthTool({
    config,
    jlincAuthDecision,
    targetAuthorized: calculatorPublic,
    targetNotAuthorized: calculatorPrivate, // Optional
  });
  const tools = [jlincAuthTool];

  const prompt = ChatPromptTemplate.fromMessages([
    ["system", "You are a helpful assistant"],
    ["placeholder", "{chat_history}"],
    ["human", "{input}"],
    ["placeholder", "{agent_scratchpad}"],
  ]);

  const agent = createToolCallingAgent({ llm, tools, prompt });

  const agentExecutor = new AgentExecutor({
    agent,
    tools,
  });

  try {
    const r = await agentExecutor.invoke({ input: "Add 1 + 1. If a function call is used, tell me the output of the function call." }, { callbacks: [tracer] });

    // The next invocation requires a reauth for any future calls to the agent:
    auth.action.name = "write";
    jlincAuthDecision.evaluate(auth);

    console.log(`\nResult`)
    console.log(`---------------------------------------------`)
    console.log(r)
  } catch (err) {
    console.error("Error calling LLM:", err);
  } finally {
    await awaitAllCallbacks();
  }
}

main()