package lc.irc.bot;

import com.theokanning.openai.completion.chat.ChatCompletionRequest;
import com.theokanning.openai.completion.chat.ChatMessage;
import com.theokanning.openai.completion.chat.ChatMessageRole;
import com.theokanning.openai.service.OpenAiService;
import lc.mecha.log.MechaLogger;
import lc.mecha.log.MechaLoggerFactory;
import org.apache.commons.lang.StringUtils;
import org.pircbotx.Configuration;
import org.pircbotx.PircBotX;
import org.pircbotx.exception.IrcException;
import org.pircbotx.hooks.ListenerAdapter;
import org.pircbotx.hooks.types.GenericMessageEvent;

import java.io.IOException;
import java.time.Duration;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.Locale;
import java.util.Random;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * Implementation of an IRC bot that implements with OpenAI to run text prompts through a model,
 * returning a response to IRC.
 *
 * @author C. Alexander Leigh
 */
public class IrcBot extends ListenerAdapter implements Runnable {
    private static final MechaLogger logger = MechaLoggerFactory.getLogger(IrcBot.class);
    public static final int MAX_HISTORY_ITEMS = 4;
    private final PircBotX bot;

    private final OpenAiService service;
    private final String botName = System.getenv("BOT_NAME");
    private final String botLetter = System.getenv("BOT_LETTER");
    private final String model = System.getenv("MODEL");

    private static final Duration DEFAULT_TIMEOUT = Duration.ofSeconds(60);
    private int maxLength = 400;
    private final Pattern p = Pattern.compile("\\G\\s*(.{1," + maxLength + "})(?=\\s|$)", Pattern.DOTALL);
    private final Random r = new Random();

    String hedgePrompt = null;
    private final HashMap<String, LinkedList<ChatMessage>> history = new HashMap<>();

    public String prePrompt = null;

    public IrcBot(String token) {
        Configuration configuration = new Configuration.Builder()
                .setName(botName) //Set the nick of the bot. CHANGE IN YOUR CODE
                .addServer("irc0.gdn.leigh-co.com", 6667)
                .addAutoJoinChannel("#leighco").addAutoJoinChannel("#gnomes")
                .addListener(this)
                .setName(botName)
                .setLogin(botName)
                .buildConfiguration();

        bot = new PircBotX(configuration);
        service = new OpenAiService(token, DEFAULT_TIMEOUT);

        switch (botName) {
            case "Werner":
                prePrompt = "Pretending to be Werner Herzog, do not mention you are pretending, do not mention OpenAI, respond to the following: ";
                hedgePrompt = "Pretending to be Werner Herzog, do not mention you are pretending, do not mention OpenAI, you just overheard the following, say something interesting about it: ";
                break;
            case "Teach":
                prePrompt = "Pretending to the dread pirate Blackbeard, do not mention you are pretending, do not mention OpenAI, respond to the following: ";
                hedgePrompt = "Pretending to be the dread pirate Blackbeard, do not mention you are pretending, do not mention OpenAI, you just overheard the following, say something interesting about it: ";
                break;
            case "Watts":
                prePrompt = "Pretending to be the philosopher Alan Watts, do not mention you are pretending, do not mention OpenAI, respond to the following: ";
                hedgePrompt = "Pretending to be the philosopher Alan Watts, do not mention you are pretending, do not mention OpenAI, you just overheard the following, say something interesting about it: ";
                break;
            case "Aurelius":
                prePrompt = "Pretending to be the Marcus Aurelius, do not mention you are pretending, do not mention OpenAI, respond to the following: ";
                hedgePrompt = "Pretending to be the Marcus Aurelius, do not mention you are pretending, do not mention OpenAI, you just overheard the following, say something interesting about it: ";
                break;
        }
    }

    @Override
    public void run() {
        try {
            bot.startBot();
        } catch (IOException | IrcException e) {
            throw new RuntimeException(e);
        }
    }

    public synchronized LinkedList<ChatMessage> getHistory(String nick) {
        LinkedList<ChatMessage> msgs = history.get(nick);
        if (msgs == null) {
            msgs = new LinkedList<>();
            history.put(nick, msgs);
        }
        return msgs;
    }

    public void push(LinkedList<ChatMessage> msgs, ChatMessage message) {
        msgs.add(message);
        logger.info("Pushing message to history. [sz: {}]", msgs.size());
        if (msgs.size() > MAX_HISTORY_ITEMS) {
            msgs.removeFirst();
        }
    }

    public void onGenericMessage(GenericMessageEvent event) {
        String msg = event.getMessage();
        String lMsg = msg.toLowerCase(Locale.ROOT);

        String prompt = null;

        if (lMsg.startsWith(botLetter.toLowerCase() + ": ")
                || lMsg.startsWith(botName.toLowerCase() + ": ")) {
            prompt = msg.substring(2);
            if (prePrompt != null) {
                prompt = prePrompt + msg.substring(2);
            }
        }

        // if (hedgePrompt != null && prompt == null && r.nextInt(300) == 0) {
        //     prompt = hedgePrompt + msg.substring(2);
        // }

        if (prompt == null) return;

        ChatMessage userMessage = new ChatMessage(ChatMessageRole.USER.value(), prompt);

        LinkedList<ChatMessage> msgs = getHistory(event.getUser().getNick());

        push(msgs, userMessage);

        ChatCompletionRequest completionRequest = ChatCompletionRequest.builder()
                .messages(msgs)
                .model(model)
                .stream(false)
                .maxTokens(500)
                .build();

        ChatMessage responseMessage = service.createChatCompletion(completionRequest).getChoices().get(0).getMessage();

        if (responseMessage == null) {
            event.respond("Hrm, I really don't know.");
            return;
        }

        push(msgs, responseMessage);

        logger.info("Result: {}", responseMessage);

        logger.info("CompletionChoice: {}", responseMessage);
        String txt = responseMessage.getContent();
        String[] lines = txt.split("\\R");
        for (String s : lines) {
            s = s.trim();
            if (!StringUtils.isBlank(s)) {
                Matcher m = p.matcher(s);
                while (m.find()) {
                    String subStr = m.group(1);
                    if (!StringUtils.isBlank(subStr))
                        event.respond(m.group(1));
                }
            }
        }
    }
}