import sys
from io import StringIO

import streamlit as st              # pip install streamlit
from code_editor import code_editor # pip install streamlit_code_editor
import ollama as ol                 # pip install ollama

st.set_page_config(layout='wide')
st.title('`Offline code completion`')

def auto_complete(model='codellama:13b-python'):
    sys_message = 'You are an AI code completion system. Generate code to complete the given Python code.'
    if 'in_code' not in st.session_state: st.session_state.in_code = ''
    res = code_editor(st.session_state.in_code, keybindings='vim')
    if len(res['id']) != 0 and (res['type'] == 'submit' or res['type'] == 'selection'):
        code = res['text']
        if code:
            ret = ol.generate(model=model, prompt=code, system=sys_message)['response']
            st.session_state.in_code = f'{code}{ret}'
            st.code(st.session_state.in_code)
        else: st.session_state.in_code = ''

def run_code(code):
    old_stdout = sys.stdout
    redirected_output = sys.stdout = StringIO()
    exec(code)
    sys.stdout = old_stdout
    st.code(redirected_output.getvalue())

auto_complete()
if st.button('RUN'):
   run_code(st.session_state.in_code)