# LLM chat for T-Deck CicruitPython version
# (also can be run on MicroPython and usual CPython)
# On the T-Deck, requires to be run from tdeck_repl for the
# keyboard to work correctly
# Usage: from llmchat import llmchat
# Created by Luxferre in 2025, released into public domain

from deck.chat import DeckChat
from deck.llm import LLMChat
from app.ed import edbuf

CONFIG_FILE = 'llmcfg.json'

# chat handlers

def send_msg(msg, state):
  print("Please wait for the answer...")
  status, ans = state.get_completion(msg)
  return ans, state

def set_sys_prompt(prompt, state):
  if prompt:
    state.config['system_prompt'] = prompt
    state.save_config()
    state.reset_context()
    return 'System prompt updated', state
  else:
    return f"System: {state.config['system_prompt']}", state

def set_temp(temp, state):
  if temp:
    state.config['temperature'] = float(temp)
    state.save_config()
    return 'Temperature updated', state
  else:
    return f"Temperature: {state.config['temperature']}", state

def set_prov(prov, state):
  if prov and prov in state.config.get('providers', {}):
    state.config['active_provider'] = prov
    state.setup_provider()
    state.save_config()
    state.reset_context()
    return 'Provider updated', state
  else:
    return f"Active provider: {state.config['active_provider']}", state

def set_model(model, state):
  if model:
    state.provider['model'] = model
    state.save_config()
    return 'Model updated', state
  else:
    return f"Active model: {state.provider.get('model')}", state

def list_models(stub, state):
  status, mlist = state.list_models()
  if status:
    resp = f"Model list for {state.config['active_provider']}:"
    for m in mlist:
      resp += f'\n- {m}'
  else:
    resp = mlist
  return resp, state

def list_providers(stub, state):
  resp = 'Provider list:'
  for p in state.config.get('providers', {}):
    resp += f'\n - {p}'
  return resp, state

def save_conversation(fname, state):
  resp = ''
  status, cnv = state.export()
  try:
    with open(fname, 'w') as f:
      f.write(cnv)
      resp = 'Conversation saved'
  except:
    resp = 'Error saving conversation'
  return resp, state

def save_last_msg(fname, state):
  resp = ''
  status, cnv = state.get_last_response()
  try:
    with open(fname, 'w') as f:
      f.write(cnv)
      resp = 'Response saved'
  except:
    resp = 'Error saving response'
  return resp, state

def extract_code_block(text, start_marker="```", end_marker="```"):
  extracted_parts = []
  search_index = 0
  while True:
    block_start = text.find(start_marker, search_index)
    if block_start == -1:
      break
    content_start = block_start + len(start_marker)
    if start_marker == "```":
      newline_after_start = text.find("\n", content_start)
      if newline_after_start != -1:
        potential_end = text.find(end_marker, content_start)
        if potential_end == -1 or newline_after_start < potential_end:
          content_start = newline_after_start + 1
    block_end = text.find(end_marker, content_start)
    if block_end == -1:
      break
    extracted_block = text[content_start:block_end].strip()
    if extracted_block:
      extracted_parts.append(extracted_block)
    search_index = block_end + len(end_marker)
  return "\n\n".join(extracted_parts)

def save_last_msg_code(fname, state):
  resp = ''
  status, cnv = state.get_last_response()
  cnv = extract_code_block(cnv)
  try:
    with open(fname, 'w') as f:
      f.write(cnv)
      resp = 'Code saved'
  except:
    resp = 'Error saving code'
  return resp, state

def edit_last_msg_code(fname, state):
  resp = ''
  status, cnv = state.get_last_response()
  cnv = extract_code_block(cnv)
  try:
    edbuf(cnv)
  except:
    resp = 'Error extracting code'
  return resp, state

def add_file(fname, state):
  status, resp = state.add_file(fname)
  return resp, state

def display_help(stub, state):
  resp = "/clear, /system, /temp, /prov, /model, /modellist, /provlist, /add, /saveconv, /savelast, /savecode, /edcode, /exit, /quit"
  return resp, state

def clear_ctx(stub, state):
  state.reset_context()
  return 'Context reset', state

# main initialization
def llmchat():
  llm = LLMChat(config_file=CONFIG_FILE, message_limit=20)
  chat = DeckChat(start_state=llm, chat_prefix='> ')
  chat.command('default', send_msg)
  chat.command('clear', clear_ctx)
  chat.command('system', set_sys_prompt)
  chat.command('temp', set_temp)
  chat.command('prov', set_prov)
  chat.command('model', set_model)
  chat.command('modellist', list_models)
  chat.command('provlist', list_providers)
  chat.command('add', add_file)
  chat.command('saveconv', save_conversation)
  chat.command('savelast', save_last_msg)
  chat.command('savecode', save_last_msg_code)
  chat.command('edcode', edit_last_msg_code)
  chat.command('help', display_help)
  chat.command('exit', lambda m,s: ('Exiting...', None))
  chat.command('quit', lambda m,s: ('Exiting...', None))
  chat.start()