import toml
from mastodon import Mastodon
import time
import html2text
from ollama import Client

# Load configuration from config.toml
with open('config.toml', 'r') as f:
    config = toml.load(f)

mastodon_config = config['mastodon']
ollama_config = config['ollama']
preamble = config.get('preamble', '')

# Initialize Mastodon client
mastodon = Mastodon(
    access_token=mastodon_config['access_token'],
    api_base_url=mastodon_config['host']
)

# Function to send text to Ollama and get a response
def ask_ollama(prompt):
    client = Client(
        host=f"http://{ollama_config['host']}:{ollama_config['port']}"
    )
    try:
        response = client.generate(f"{ollama_config['model']}", f"{preamble}\n{prompt}")
        print(f"{ollama_config['model']}", f"{preamble}\n{prompt}")
        return response['response']
    except Exception as e:
        print(f"HTTP error occurred: {e}")
        return "Error contacting Ollama server."

# Function to process mentions
def process_mentions():
    since_id = None
    while True:
        mentions = mastodon.notifications(since_id=since_id) if since_id else mastodon.notifications()
        for mention in mentions:
            if mention['type'] == 'mention':
                content = html2text.html2text(mention['status']['content'])
                user = mention['status']['account']['acct']
                status_id = mention['status']['id']
                print(f"Mention from {user}: {content}")

                response_text = ask_ollama(content)
                print(f"Ollama response: {response_text}")

                mastodon.status_post(
                    status=f"@{user} {response_text}",
                    in_reply_to_id=status_id,
                    visibility=mastodon_config.get('default_visibility', 'public')
                )
                since_id = mention['id']

        time.sleep(30)

if __name__ == "__main__":
    process_mentions()