# deck.llm: a class with (online) LLM interactions for T-Deck
# Supports external config files and multiple providers
# Created by Luxferre in 2025, released into public domain

import json, re
try:
  from deck.time import time
except:
  import time

# 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()
  if status:
    print("We're online, initing requests lib...")
    requests = net.init_requests()
except:
  import requests

class LLMChat:
  """
  Class for chat-like LLM interactions
  """
  def __init__(self, config_file='llmcfg.json', message_limit=10):
    self.config = self.load_config(config_file)
    self.config_file = config_file
    self.message_limit = message_limit
    self.messages = []
    self.setup_provider()
    self.reset_context()

  def load_config(self, config_file='llmcfg.json'):
    """
    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(self.config_file, 'w') as f:
        json.dump(self.config, f)
      return (True, "Configuration saved.")
    except OSError as e:
      return (False, 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:
      return (False, f"Provider {active_name} not found.")
        
    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)
    return (True, "Provider ready")

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

  def guarded_post(self, url, headers={}, payload={}):
    """
    Performs a POST request trying to bypass possible 429 errors with heuristics.
    """
    response = requests.post(url, headers=headers, json=payload)
    if response.status_code == 429:
      restext = response.text
      try:
        resdata = response.json()
        restext = data['error']['message']
      except:
        pass
      response.close()
      # try and parse something like "please wait for 0.1 seconds"
      sleeptime = 1.0
      try:
        sleeptime = float(re.search(r'\d+\.\d+', restext).group(0))
      except:
        pass
      time.sleep(sleeptime + 0.1)
      return self.guarded_post(url, headers, payload)
    else:
      data = response.json()
      response.close()
      return data

  def get_completion(self, user_input):
    """
    Sends request to LLM.
    """
    self.messages.append({"role": "user", "content": user_input})
    
    if len(self.messages) > (self.message_limit + 1):
      self.messages = [self.messages[0]] + self.messages[-self.message_limit:]

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

    url = f"{self.provider.get('base_url').rstrip('/')}/chat/completions"
 
    try:
      data = self.guarded_post(url, self.headers, payload)
      bot_content = data['choices'][0]['message']['content']
      self.messages.append({"role": "assistant", "content": bot_content})
      return (True, bot_content)
    except Exception as e:
      if len(self.messages) > 0: self.messages.pop()
      return (False, f"Error: {e}")

  def get_last_response(self):
    """
    Get the most recent assistant's response.
    """
    idx = len(self.messages) - 1
    while idx > 0 and self.messages[idx]['role'] != 'assistant':
      idx -= 1
    if idx == 0:
      return (False, 'No assistant messages detected')
    else:
      return (True, self.messages[idx]['content'])

  def export(self):
    """
    Export the entire conversation as a string.
    """
    out = ''
    for msg in self.messages:
      out += f'{msg['role']}: {msg['content'].strip()}\n'
    return (True, out.strip())

  def add_file(self, fname):
    """
    Add a (text) file to the context.
    """
    try:
      with open(fname, 'r') as f:
        text = f.read()
        prompt = f'Contents of the file {fname}:\n\n{text}'
        self.messages.append({'role': 'user', 'content': prompt})
        return (True, 'File content added')
    except:
      return (False, 'Could not read the file')

  def list_models(self):
    """
    Lists available models for the active provider.
    """
    url = f"{self.provider.get('base_url').rstrip('/')}/models"
    try:
      response = requests.get(url, headers=self.headers)
      data = response.json()
      response.close()
      res = []
      if 'data' in data:
        res = [m.get('id') for m in data['data']]
      return True, res
    except Exception as e:
      return False, f"Error: {e}"