#!/usr/bin/env python3

import os
import argparse

import yaml

from ztl.core.server import TaskServer
from ztl.core.client import RemoteTask
from ztl.core.protocol import State, Task
from ztl.core.task import ExecutableTask, TaskExecutor, TaskController

from ollama import Client


class OLLamaTask(ExecutableTask):

  def __init__(self, client, model, content, history, remotes, component):
    self.client = client
    self.model = model
    self.content = content
    self.history = history
    self.remotes = remotes
    self.robot = None
    if len(component.split("/")) == 2:
        self.robot, self.component = component.split("/")
        

  def execute(self):
    print("OLLAMA REQUEST >>> %s" % self.content)
    message = {
          'role': 'user',
          'content': self.content,
          'stream': False
        }

    self.history.append(message)

    response = self.client.chat(
      model=self.model,
      messages=self.history
    )

    self.history.append(response.message)
    print("OLLAMA RESPONSE <<< %s" % response.message.content)

    if self.robot is None:
      return response.message.content
    
    if not self.robot in self.remotes:
      raise RuntimeError("No remote for handler '%s', exiting after prompt." % self.robot)

    utterance = response.message.content.encode('ascii', errors="ignore").decode()
    mid, reply = self.remotes[self.robot].trigger(Task.encode(self.robot, self.component, utterance))
    if mid == -1:
      raise RuntimeError("Failed to trigger robot utterance: %s" % reply)

    status, reply = self.remotes[self.robot].wait(mid, timeout=15)
    if not status == State.COMPLETED:
      raise RuntimeError("Failed to complete robot utterance: %s" %reply)

    print("ROBOT RESPONSE  <<< %s" % reply)
    return reply

  def abort(self):
    return False


class PromptController(TaskController):

  def __init__(self, configfile):
    super(PromptController, self).__init__()
    self.histories = {}
    self.clients = {}
    self.tasks = {}
    self.remotes = {}
    
    with open(configfile) as f:
      self.config = yaml.safe_load(f)

      rs = self.config["remotes"]
      for r in rs.keys():
        print("Initialising remote task interface '%s'..." % r)
        self.remotes[r] = RemoteTask(rs[r]["host"], rs[r]["port"], rs[r]["scope"])

  def add_handler(self, name, cls, client):
    self.histories[name] = []
    self.clients[name] = client
    self.tasks[name] = cls

  def assign(self, handler, component, goal):

    ai, model = handler.split("/")

    if ai in self.clients:
      return self.tasks[ai], self.clients[ai], model, goal, self.histories[ai], self.remotes, component
    else:
      raise RuntimeError("No such handler '%s'" % handler)


def main_cli():
  
  cfg_file = os.environ.get('XDG_CONFIG_HOME', os.environ.get('HOME', '/home/demo') + '/.config') + '/zmq-remotes.yaml'

  parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
  parser.add_argument("-c", "--config", type=str,
                      help="Configuration file location.", default=cfg_file)
  parser.add_argument("-p", "--port", type=int,
                      help="The port to listen listen on.", default=5571)
  parser.add_argument("-s", "--scope", type=str,
                      help="The scope that to respond to.", default="/talk")
  parser.add_argument("-l", "--ollama", type=str,
                      help="The OLLama instance address.", default=os.environ.get('OLLAMA_HOST', 'http://localhost:11434'))


  args, unknown = parser.parse_known_args()

  # client = RemoteTask(args.robot_host, args.robot_port, args.robot_scope)

  pc = PromptController(args.config)
  pc.add_handler("ollama", OLLamaTask, Client(host=args.ollama))

  server = TaskServer(args.port)
  server.register(args.scope, pc)
  server.listen()

if __name__ == "__main__":

  main_cli()
