import { BaseHandler } from "./base-handler.ts"
import { Bot, Channel, Message } from "../snek/snek-socket.ts"
import { trim, trimStart } from "npm:lodash-es"
import {
  ChatSessionModelFunctions,
  ChatWrapper,
  defineChatSessionFunction,
  GemmaChatWrapper,
  GeneralChatWrapper,
  getLlama,
  LLamaChatPromptOptions,
  LlamaChatSession,
  LlamaContext,
  LlamaModel,
  resolveChatWrapper,
  Token,
} from "npm:node-llama-cpp"
import { getLogger } from "@logtape/logtape"
import { deepMerge } from "@std/collections/deep-merge"
import { fetchRemoteFunctions } from "../ai/remote-functions.ts"

const llama = await getLlama()

const textEncoder = new TextEncoder()

function printSync(input: string | Uint8Array, to = Deno.stdout) {
  let bytesWritten = 0
  const bytes = typeof input === "string" ? textEncoder.encode(input) : input
  while (bytesWritten < bytes.length) {
    bytesWritten += to.writeSync(bytes.subarray(bytesWritten))
  }
}
const logger = getLogger(["llama-gen-handler"])

const optionsGenerator = <
  const Functions extends ChatSessionModelFunctions | undefined =
    | ChatSessionModelFunctions
    | undefined,
  LLamaOptions extends LLamaChatPromptOptions<Functions> =
    LLamaChatPromptOptions<Functions>,
>(
  model: LlamaModel,
  debugOutput: boolean = true,
  defaultTimeout = 5 * 60 * 1000,
  options?: LLamaOptions,
): LLamaOptions => {
  const manager = AbortSignal.timeout(defaultTimeout)

  const defaultOptions = {
    repeatPenalty: {
      lastTokens: 24,
      penalty: 1.12,
      penalizeNewLine: true,
      frequencyPenalty: 0.02,
      presencePenalty: 0.02,
      punishTokensFilter: (tokens: Token[]) => {
        return tokens.filter((token) => {
          const text = model.detokenize([token])

          // allow the model to repeat tokens
          // that contain the word "better"
          return !text.toLowerCase().includes("@")
          // TODO: Exclude usernames
        })
      },
    },
    temperature: 0.7,
    minP: 0.03,

    // topK: 64,
    // topP: 0.95,
    // minP: 0.01,
    signal: manager,
    stopOnAbortSignal: true,
  } as LLamaOptions

  if (debugOutput) {
    defaultOptions.onResponseChunk = (chunk) => {
      options?.onResponseChunk?.(chunk)
      const isThoughtSegment = chunk.type === "segment" &&
        chunk.segmentType === "thought"

      if (
        chunk.type === "segment" && chunk.segmentStartTime != null
      ) {
        printSync(` [segment start: ${chunk.segmentType}] `)
      }

      printSync(chunk.text)

      if (chunk.type === "segment" && chunk.segmentEndTime != null) {
        printSync(` [segment end: ${chunk.segmentType}] `)
      }
    }
  }

  return deepMerge(defaultOptions, options ?? {})
}

export class LLamaFuncHandler extends BaseHandler {
  joinMode = new Map<string, boolean>()
  streamMode = false
  debugLogResponses = true

  systemPrompt: string

  #activeModel: string

  #model: LlamaModel | null = null
  #context: LlamaContext | null = null

  #chatWrapper: ChatWrapper | null = null

  #session: LlamaChatSession | null = null
  #subCommands = {
    "prompt": this.prompt.bind(this),
    "join": this.join.bind(this),
    "stream": this.stream.bind(this),
    "reset": this.reset.bind(this),
  } as Record<string, (command: string, message: Message, bot: Bot) => void>

  constructor(
    activeModel: string,
    systemPrompt: string = "You are an AI chatbot.",
  ) {
    super("")
    this.#activeModel = activeModel
    this.systemPrompt = systemPrompt
    this.autoTyping = false
  }

  async calculateSystemPrompt(): Promise<string> {
    return this.systemPrompt
  }

  override async bind(bot: Bot): Promise<void> {
    await super.bind(bot)
    this.prefix = this.user!.username.toLowerCase()

    this.#model = await llama.loadModel({
      modelPath: this.#activeModel,
      defaultContextFlashAttention: true,
    })

    this.#context = await this.#model.createContext({
      flashAttention: true,
    })

    logger.info("Model loaded", {
      batchSize: this.#context.batchSize,
      contextSize: this.#context.contextSize,
    })

    this.#chatWrapper = //new Llama3ChatWrapper()
      resolveChatWrapper(this.#model) ?? new GeneralChatWrapper()

    this.#session = new LlamaChatSession({
      contextSequence: this.#context.getSequence(),
      chatWrapper: this.#chatWrapper,
      systemPrompt: await this.calculateSystemPrompt(),
    })

    // const channels = await bot.channels

    // const channel = channels.find((v) => v.tag === "public") || channels[0]

    // if (channel) {
    // await bot.sendMessage(
    //   channel.uid,
    //   await this.cleanResponse(
    //     await this.#session.prompt(
    //       "Welcome to chat, greet everyone\n",
    //       optionsGenerator(this.#model, this.debugLogResponses),
    //     ),
    //     null,
    //     bot,
    //   ),
    // )
    // this.#session.resetChatHistory()
    // }

    // logger.info("LLamaHandler bound to bot")
  }

  cachedChannels: Channel[] | null = null
  async isInJoinMode(channelUID: string, bot: Bot): Promise<boolean> {
    if (!this.joinMode.has(channelUID)) {
      if (!this.cachedChannels) {
        this.cachedChannels = await bot.channels
      }
      const channel = this.cachedChannels?.find((c) => c.uid === channelUID)
      if (channel) {
        this.joinMode.set(channelUID, channel.tag === "dm")
      } else {
        logger.warn("Channel not found in cached channels", { channelUID })
        this.cachedChannels = await bot.channels
      }
    }
    return this.joinMode.get(channelUID) ?? false
  }

  override async isMatch(message: Message): Promise<boolean> {
    return message
          .userUID !== this.user?.uid &&
      message.isFinal &&
      (await this.isInJoinMode(message.channelUID, message.bot) ||
        trim(message?.message, " `").toLowerCase().includes(this.prefix))
  }

  async cleanResponse(
    response: string,
    message: Message | null,
    bot: Bot,
  ): Promise<string> {
    const session = this.#session
    const user = this.user

    response = trim(response.replace(/<think>.*?<\/think>/gs, ""), '" \t')
    let lwResponse = response.toLowerCase()

    if (lwResponse.startsWith("ai")) {
      response = response.substring(2).trim()
      lwResponse = response.toLowerCase()
    }

    if (user && lwResponse.startsWith(`@${user.username.toLowerCase()}:`)) {
      response = response.substring(user.username.length + 2).trim()
      lwResponse = response.toLowerCase()
    }

    if (
      message && lwResponse.startsWith(`@${message.username.toLowerCase()}:`)
    ) {
      response = response.substring(message.username.length + 2).trim()
      lwResponse = response.toLowerCase()
    }

    if (user && lwResponse.startsWith(`${user.username.toLowerCase()}:`)) {
      response = response.substring(user.username.length + 2).trim()
      lwResponse = response.toLowerCase()
    }

    if (
      message && lwResponse.startsWith(`${message.username.toLowerCase()}:`)
    ) {
      response = response.substring(message.username.length + 2).trim()
      lwResponse = response.toLowerCase()
    }

    response = trimStart(response, ":").trim()
    response = trim(response, '"')
    return response
  }

  async join(command: string, message: Message, bot: Bot): Promise<void> {
    this.joinMode.set(
      message.channelUID,
      !await this.isInJoinMode(message.channelUID, bot),
    )
  }

  async stream(command: string, message: Message, bot: Bot): Promise<void> {
    this.streamMode = !this.streamMode
  }

  async reset(command: string, message: Message, bot: Bot): Promise<void> {
    const session = this.#session
    if (!session) {
      return
    }

    await session.resetChatHistory()
    bot.sendMessage(
      message.channelUID,
      await this.cleanResponse(
        await session.prompt(
          "Your memory was just reset. Welcome to chat, greet everyone\n",
          optionsGenerator(this.#model!, this.debugLogResponses),
        ),
        message,
        bot,
      ),
    )
    await session.resetChatHistory()
  }

  async prompt(command: string, message: Message, bot: Bot): Promise<void> {
    const session = this.#session
    const user = this.user
    if (!session || !user) {
      return
    }

    let msgSoFar = ""

    let sentMessageInfo: Promise<unknown> | null = null

    let streamId: number | undefined = undefined
    if (this.streamMode) {
      streamId = setInterval(async () => {
        if (
          msgSoFar.length < 1 ||
          (msgSoFar.startsWith("@") && !msgSoFar.includes(" "))
        ) {
          return
        }
        sentMessageInfo = sentMessageInfo?.then(async (msgInfo) => {
          const msg = await this.cleanResponse(msgSoFar, message, bot)
          try {
            const msgRes = await bot.updateMessageText(
              msgInfo as string,
              msg,
            )
            if ("error" in msgRes) {
              console.error("Error updating message text", msgRes)
              msgSoFar = ""
              return bot.sendMessage(message.channelUID, msg)
            }
          } catch (data) {
            console.error(data)
          }
          return msgInfo
        }) ??
          bot.sendMessage(
            message.channelUID,
            await this.cleanResponse(msgSoFar, message, bot),
          )
      }, 50)
    }

    console.log("Prompting model", {
        model: this.#activeModel,
        command,
        message,
        user: user.username,
    })
    let response = await session.prompt(
      `@${message.username}: ${message.message}`,
      optionsGenerator(this.#model!, this.debugLogResponses, 5 * 60 * 1000, {
        onTextChunk: (text) => {
          bot.setTyping(message.channelUID).catch(() => {})

          msgSoFar += text

          if (this.streamMode) {
            // sentMessageInfo = sentMessageInfo
            //   ? sentMessageInfo.then(([msgInfo, msg]) => {
            //     const newMsg = msg + text
            //     return this.cleanResponse(newMsg, message, bot).then((msg) => {
            //       return bot.updateMessageText(
            //         msgInfo,
            //         msg,
            //       ).catch(console.error).then(() => [msgInfo, newMsg])
            //     })
            //   })
            //   : bot.sendMessage(
            //     message.channelUID,
            //     text,
            //   ).then((msg) => [msg as string, text])
          }
        },
        functions: {
          generateImage: defineChatSessionFunction({
            description: "Generate an image from a prompt",
            params: {
              type: "object",
              properties: {
                prompt: {
                  type: "string",
                  description: "The prompt to generate the image from",
                },
              },
              required: ["prompt"],
            },
            handler: async ({ prompt }) => {
              bot.sendMessage(message.channelUID, "@abot prompt " + prompt)
            },
          }),
          now: defineChatSessionFunction({
            description: "Get the current time",
            params: {
              type: "object",
              properties: {},
            },
            handler: async () => {
              const now = new Date()
              return now.toLocaleString()
            },
          }),
          ...(await fetchRemoteFunctions().catch((e) => {
            logger.error("Failed to fetch remote functions", { error: e })
            return ({})
          })),
        },
      }),
    )

    response = await this.cleanResponse(response, message, bot)

    clearInterval(streamId)

    if (sentMessageInfo) {
      const msgInfo = await sentMessageInfo

      bot.updateMessageText(
        msgInfo,
        response,
      )
    } else {
      bot.sendMessage(message.channelUID, response)
    }
  }

  override async handleMessage(message: Message, bot: Bot): Promise<void> {
    const user = this.user
    if (!user) {
      return
    }

    if (
      message.message.toLowerCase().startsWith(
        `@${user.username.toLowerCase()}`,
      )
    ) {
      const newMessage = message.message.substring(
        message.message.indexOf(" "),
        message.message.length,
      ).trim()
      if (newMessage && message.message.includes(" ")) {
        const [[_, command, rest]] = newMessage.matchAll(/^(\S+)\s*(.*)$/gs)

        if (command in this.#subCommands) {
          return this.#subCommands[command]?.(rest.trim(), message, bot)
        }
      }
    }

    return this.prompt(message.message, message, bot)
  }
}