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

import sys, os, json

# init the tdeck_repl input
try:
  from tdeck_repl import input
except:
  pass

# init requests library in a platform-agnostic way
requests = None
try:
  from deck import net
  print('Waiting for online status...')
  status, stext, myip = net.wifi_connect_wait()
  print(status, stext, myip)
  if status:
    print("We're online, initing requests lib...")
    requests = net.init_requests()
except:
  import requests

CONFIG_FILE = 'llmcfg.json'

class LLMChat:
  def __init__(self):
    self.config = self.load_config()
    self.messages = []
    self.setup_provider()
    self.reset_context()

  def load_config(self):
    """Loads config from JSON file with CircuitPython fallback."""
    try:
      with open(CONFIG_FILE, 'r') as f:
        return json.load(f)
    except Exception as e:
      print(f"Error: Could not load {CONFIG_FILE} ({e})")
      # Minimal fallback for embedded if file is missing
      return {
        "system_prompt": "Concise assistant",
        "temperature": 0.7,
        "active_provider": "local",
        "providers": {"pollinations": {"base_url": "https://text.pollinations.ai/openai", "api_key": "", "model": "openai-fast"}}
      }

  def save_config(self):
    """Saves config. Note: May fail on CircuitPython if USB is connected."""
    try:
      with open(CONFIG_FILE, 'w') as f:
        json.dump(self.config, f)
      print("Configuration saved.")
    except OSError as e:
      print(f"Save failed: {e}. (Is filesystem read-only?)")

  def setup_provider(self):
    """Sets up headers and provider settings."""
    active_name = self.config.get('active_provider')
    providers = self.config.get('providers', {})
    
    if active_name not in providers:
      print(f"Provider {active_name} not found.")
      return
        
    self.provider = providers[active_name]
    self.headers = {
      "Content-Type": "application/json",
      "Authorization": f"Bearer {self.provider.get('api_key', '')}"
    }
    extra = self.provider.get('extra_headers', {})
    if extra: self.headers.update(extra)

  def reset_context(self):
    """Clears history and re-adds system prompt."""
    self.messages = [{"role": "system", "content": self.config.get('system_prompt', '')}]
    print("Context cleared.")

  def paginate_output(self, text, lines_per_page=8):
    """Splits output for small terminals."""
    lines = text.split('\n')
    for i in range(0, len(lines), lines_per_page):
      chunk = lines[i:i + lines_per_page]
      for line in chunk: print(line)
      if i + lines_per_page < len(lines):
        input("-- Press Enter --")

  def get_completion(self, user_input):
    """Sends request to LLM."""
    self.messages.append({"role": "user", "content": user_input})
    
    # Memory safety for embedded: keep only last 10 messages
    if len(self.messages) > 11:
      self.messages = [self.messages[0]] + self.messages[-10:]

    payload = {
      "model": self.provider.get('model'),
      "messages": self.messages,
      "temperature": self.config.get('temperature', 0.7)
    }

    url = f"{self.provider.get('base_url').rstrip('/')}/chat/completions"
    print("Please wait for the answer...")
    
    try:
      response = requests.post(url, headers=self.headers, json=payload)
      data = response.json()
      response.close() # Important for CircuitPython memory
      
      bot_content = data['choices'][0]['message']['content']
      self.messages.append({"role": "assistant", "content": bot_content})
      return bot_content
    except Exception as e:
      if len(self.messages) > 0: self.messages.pop()
      return f"Error: {e}"

  def list_models(self):
    """Lists available models."""
    url = f"{self.provider.get('base_url').rstrip('/')}/models"
    try:
      response = requests.get(url, headers=self.headers)
      data = response.json()
      response.close()
      print(f"Models for {self.config['active_provider']}:")
      if 'data' in data:
        for m in data['data']: print(f" - {m.get('id')}")
    except Exception as e:
      print(f"Error: {e}")

  def run(self):
    print(f"LLM CLI | {self.config.get('active_provider')} | {self.provider.get('model')}")
    
    while True:
      try:
        user_input = input("\n>> ").strip()
        if not user_input: continue

        if user_input.startswith('/'):
          parts = user_input.split(maxsplit=1)
          cmd, arg = parts[0].lower(), parts[1] if len(parts) > 1 else None

          if cmd in ['/exit', '/quit']: break
          elif cmd == '/clear': self.reset_context()
          elif cmd == '/system':
            if arg:
              self.config['system_prompt'] = arg
              self.save_config()
              self.reset_context()
            else: print(f"System: {self.config['system_prompt']}")
          elif cmd == '/temp':
            if arg:
              self.config['temperature'] = float(arg)
              self.save_config()
            else: print(f"Temp: {self.config['temperature']}")
          elif cmd == '/prov':
            if arg and arg in self.config.get('providers', {}):
              self.config['active_provider'] = arg
              self.setup_provider()
              self.save_config()
              self.reset_context()
            else: print(f"Active: {self.config['active_provider']}")
          elif cmd == '/model':
            if arg:
              self.provider['model'] = arg
              self.save_config()
            else: print(f"Model: {self.provider.get('model')}")
          elif cmd == '/modellist': self.list_models()
          elif cmd == '/provlist':
            for p in self.config.get('providers', {}): print(f" - {p}")
          elif cmd == '/help':
            print("/clear, /system, /temp, /prov, /model, /modellist, /provlist, /exit")
          else: print("Unknown command.")
          continue

        ans = self.get_completion(user_input)
        print("")
        self.paginate_output(ans)

      except KeyboardInterrupt: break
    print("\nExiting.")

# no main, direct init
chat = LLMChat()
chat.run()