import toml
import requests
import mastodon
import os
import html2text
import ollama
import time

CONFIG_PATH = 'config.toml'
STATE_FILE = 'replied_mentions.txt'

def load_config(path):
    with open(path, 'r') as file:
        return toml.load(file)

def load_state():
    if not os.path.exists(STATE_FILE):
        return set()
    with open(STATE_FILE, 'r') as file:
        return set(line.strip() for line in file.readlines())

def save_state(state):
    with open(STATE_FILE, 'w') as file:
        for item in state:
            file.write(f"{item}\n")

def send_to_ollama(prompt, config):
    client = ollama.Client(
        host=f"http://{config['ollama']['host']}:{config['ollama']['port']}"
    )
    try:
        response = client.generate(f"{config['ollama']['model']}", f"{config['ollama']['preamble']}\n{prompt}")
        print(f"{config['ollama']['model']}", f"{config['ollama']['preamble']}\n{prompt}")
        return response['response']
    except Exception as e:
        print(f"HTTP error occurred: {e}")
        return "Error contacting Ollama server."

def main():
    config = load_config(CONFIG_PATH)
    state = load_state()

    masto = mastodon.Mastodon(
        api_base_url=config['mastodon']['host'],
        access_token=config['mastodon']['access_token']
    )

    mentions = masto.notifications(mentions_only=True)

    for mention in mentions:
        if mention['id'] in state:
            continue

        content = mention['status']['content']
        username = mention['status']['account']['username']
        prompt = html2text.html2text(content)
        
        if username != "maxheadroom":
            continue

        response = send_to_ollama(prompt, config)
        if not response or config['mastodon']['bot_username'] in response:
            continue

        masto.status_post(
            status=f"@{username} {response}",
            in_reply_to_id=mention['status']['id'],
            visibility=config['mastodon']['default_visibility']
        )

        state.add(mention['id'])

    save_state(state)

if __name__ == "__main__":
    while True:
        main()
        time.sleep(30)