import { makeCoreModuleName, makeFqName } from '../util.js';
import { Environment, makeEventEvaluator, parseAndEvaluateStatement } from '../interpreter.js';
import { fetchModule, Instance, instanceToObject, isModule } from '../module.js';
import { provider } from '../agents/registry.js';
import {
  AgentServiceProvider,
  AIResponse,
  assistantMessage,
  humanMessage,
  systemMessage,
} from '../agents/provider.js';
import { AIMessage, BaseMessage, HumanMessage } from '@langchain/core/messages';
import { PlannerInstructions } from '../agents/common.js';

export const CoreAIModuleName = makeCoreModuleName('ai');

export default `module ${CoreAIModuleName}

entity llm {
    name String @id,
    service String @default("openai"),
    config Map @optional
}

entity agent {
    name String @id,
    type @oneof("chat", "planner") @default("chat"),
    instruction String @optional,
    tools String[] @optional,
    documents String[] @optional,
    llm String
}

entity agentChatSession {
    id String @id,
    messages String
}

workflow findAgentChatSession {
  {agentChatSession {id? findAgentChatSession.id}} as [sess];
  sess
}

workflow saveAgentChatSession {
  upsert {agentChatSession {id saveAgentChatSession.id, messages saveAgentChatSession.messages}}
}
`;

export const AgentFqName = makeFqName(CoreAIModuleName, 'agent');

const ProviderDb = new Map<string, AgentServiceProvider>();

export class Agent {
  llm: string = '';
  name: string = '';
  chatId: string | undefined;
  instruction: string = '';
  type: string = 'chat';
  tools: string[] | undefined;

  private constructor() {}

  static FromInstance(agentInstance: Instance): Agent {
    return instanceToObject<Agent>(agentInstance, new Agent());
  }

  isPlanner(): boolean {
    return (this.tools && this.tools.length > 0) || this.type == 'planner';
  }

  async invoke(message: string, env: Environment) {
    const p = await findProviderForLLM(this.llm, env);
    const agentName = this.name;
    const chatId = this.chatId || agentName;
    const sess: Instance | null = await findAgentChatSession(chatId, env);
    let msgs: BaseMessage[] | undefined;
    const isplnr = this.isPlanner();
    if (sess) {
      msgs = sess.lookup('messages');
    } else {
      msgs = [systemMessage(this.instruction)];
    }
    if (msgs) {
      const sysMsg = msgs[0];
      if (isplnr) {
        const newSysMsg = systemMessage(
          `${PlannerInstructions}\n${this.toolsAsString()}\n${this.instruction}`
        );
        msgs[0] = newSysMsg;
      }
      msgs.push(humanMessage(message));
      const response: AIResponse = await p.invoke(msgs);
      msgs.push(assistantMessage(response.content));
      if (isplnr) {
        msgs[0] = sysMsg;
      }
      await saveAgentChatSession(chatId, msgs, env);
      env.setLastResult(response.content);
    } else {
      throw new Error(`failed to initialize messages for agent ${agentName}`);
    }
  }

  private toolsAsString(): string {
    if (this.tools) {
      return this.tools
        .filter((s: string) => {
          return isModule(s);
        })
        .map((moduleName: string) => {
          return fetchModule(moduleName).toString();
        })
        .join('\n');
    } else {
      return '';
    }
  }
}

export async function findAgentByName(name: string, env: Environment): Promise<Agent> {
  await parseAndEvaluateStatement(`{agentlang_ai/agent {name? "${name}"}}`, undefined, env);
  const result = env.getLastResult();
  if (result instanceof Array && result.length > 0) {
    const agentInstance: Instance = result[0];
    return Agent.FromInstance(agentInstance);
  } else {
    throw new Error(`Failed to fine agent ${name}`);
  }
}

export async function findProviderForLLM(
  llmName: string,
  env: Environment
): Promise<AgentServiceProvider> {
  let p: AgentServiceProvider | undefined = ProviderDb.get(llmName);
  if (p == undefined) {
    const result: Instance[] = await parseAndEvaluateStatement(
      `{${CoreAIModuleName}/llm {name? "${llmName}"}}`,
      undefined,
      env
    );
    if (result.length > 0) {
      const llm: Instance = result[0];
      const service = llm.lookup('service');
      const pclass = provider(service);
      const providerConfig: Map<string, any> =
        llm.lookup('config') || new Map().set('service', service);
      p = new pclass(providerConfig);
      if (p) ProviderDb.set(llmName, p);
    }
  }
  if (p) {
    return p;
  } else {
    throw new Error(`Failed to load provider for ${llmName}`);
  }
}

const evalEvent = makeEventEvaluator(CoreAIModuleName);

type GenericMessage = {
  role: 'system' | 'user' | 'assistant';
  content: string;
};

function asBaseMessages(gms: GenericMessage[]): BaseMessage[] {
  return gms.map((gm: GenericMessage): BaseMessage => {
    switch (gm.role) {
      case 'user': {
        return humanMessage(gm.content);
      }
      case 'assistant': {
        return assistantMessage(gm.content);
      }
      default: {
        return systemMessage(gm.content);
      }
    }
  });
}

function asGenericMessages(bms: BaseMessage[]): GenericMessage[] {
  return bms.map((bm: BaseMessage): GenericMessage => {
    if (bm instanceof HumanMessage) {
      return { role: 'user', content: bm.text };
    } else if (bm instanceof AIMessage) {
      return { role: 'assistant', content: bm.text };
    } else {
      return { role: 'system', content: bm.text };
    }
  });
}

export async function findAgentChatSession(
  chatId: string,
  env: Environment
): Promise<Instance | null> {
  const result: Instance | null = await evalEvent('findAgentChatSession', { id: chatId }, env);
  if (result) {
    result.attributes.set('messages', asBaseMessages(JSON.parse(result.lookup('messages'))));
  }
  return result;
}

export async function saveAgentChatSession(chatId: string, messages: any[], env: Environment) {
  await evalEvent(
    'saveAgentChatSession',
    { id: chatId, messages: JSON.stringify(asGenericMessages(messages)) },
    env
  );
}

export function agentName(agentInstance: Instance): string {
  return agentInstance.lookup('name');
}