#!/bin/python3
# Chat with an intelligent assistant in your terminal
from ollama import Client
import re
import pyperclip
import sys
import argparse

model = 'mistral-nemo'

pattern = r'```[a-z]*\n[\s\S]*?\n```'
line_pattern = r'`[a-z]*[\s\S]*?`'

def extract_code_block(markdown_text):
    # Use the regular expression pattern to find all matches in the markdown text
    matches = re.finditer(pattern, markdown_text)

    # Iterate over the matches and extract the code blocks
    code_blocks = []
    for match in matches:
        code_block = match.group(0)

        # Add the code block to the list of code blocks
        code_blocks.append('\n'.join(code_block.split('\n')[1:-1]))

    if len(code_blocks) == 0:
        line_matches = re.finditer(line_pattern, markdown_text)
        for match in line_matches:
            code_block = match.group(0)
            code_blocks.append(code_block[1:-1])
    return code_blocks

def copy_string_to_clipboard(string):
    try:
        pyperclip.copy(string)
    except:
        return

# Point to the local server
client = Client(host='hayden-desktop:11434')

code_history = []

history = [
        {"role": "system", "content": "You are a helpful, smart, kind, and efficient AI assistant. You always fulfill the user's requests accurately and concisely."},
        ]

def chat(message, stream=True):
    history.append({"role": "user", "content": message})
    completion = client.chat(
            model=model,
            messages=history,
            stream=stream
            )
    result = ''
    for chunk in completion:
        if stream:
            print(chunk['message']['content'], end='', flush=True)
        result += chunk['message']['content']
    if stream:
        print()
    history.append({"role": 'assistant', 'content': result})
    return result

def parse_args():
    # Create the parser
    parser = argparse.ArgumentParser(description='Copy and open a source file in TextEdit')
    # Add the --follow-up (-f) argument
    parser.add_argument('--follow-up', '-f', nargs='?', const=True, default=False, help='Ask a follow up question when piping in context')
    # Add the --copy (-c) argument
    parser.add_argument('--copy', '-c', action='store_true', help='copy a codeblock if it appears')
    # Add the --shell (-s) argument
    parser.add_argument('--shell', '-s', nargs='?', const=True, default=False, help='output a shell command that does as described')
    # Add the --model (-m) argument
    parser.add_argument('--model', '-m', nargs='?', const=True, default=False, help='Specify model')
    # Parse the arguments
    return parser.parse_args()

def arg_follow_up(args):
    sys.stdin = open('/dev/tty')
    if args.follow_up != True:
        second_input = args.follow_up
    else:
        second_input = input('> ')
    return second_input

def arg_shell(args):
    query = 'Form a shell command based on the following description. Only output a working shell command and then <|eot_id|>.\nDescription: '
    if args.shell != True:
        query += args.shell
    else:
        query += input('> ')
    result = chat(query, False)
    result = blocks[0] if len(blocks := extract_code_block(result)) else result
    print(result)
    copy_string_to_clipboard(result)

def handle_piped_input(args):
    all_input = sys.stdin.read()
    query = 'Use the following context to answer the question. There will be no follow up questions from the user so make sure your answer is complete:\nSTART CONTEXT\n' + all_input + '\nEND CONTEXT\n'
    if args.copy:
        query += 'Answer the question using a codeblock for any code or shell scripts\n'
    if args.follow_up:
        query += arg_follow_up(args)
    query += '\n'

    result = chat(query)
    blocks = extract_code_block(result)
    if args.copy and len(blocks):
        copy_string_to_clipboard(blocks[0])

def handle_non_piped_input(args):
    if args.shell:
        arg_shell(args)
        exit()
    if args.follow_up:
        user_input = arg_follow_up(args)
        chat(user_input)
        exit()
    while True:
        try:
            user_input = input('> ')
        except (EOFError, KeyboardInterrupt):
            print()
            exit()
        else:
            chat(user_input)

def main():
    args = parse_args()
    if args.model:
        global model 
        model = args.model
    if not sys.stdin.isatty():
        handle_piped_input(args)
    else:
        handle_non_piped_input(args)

if __name__ == '__main__':
    main()