#!/usr/bin/env python3

import os
import argparse

from ztl.core.server import TaskServer
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):
    self.client = client
    self.model = model
    self.content = content
    self.history = history

  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)
    return response.message.content

  def abort(self):
    return False


class PromptController(TaskController):

  def __init__(self):
    super().__init__(self)
    self.histories = {}
    self.clients = {}
    self.tasks = {}

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

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


def main_cli():

  parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
  parser.add_argument("-p", "--port", type=int,
                      help="The port to listen listen on.", default=5570)
  parser.add_argument("-s", "--scope", type=str,
                      help="The scope that to respond to.", default="/prompt")
  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()

  server = TaskServer(args.port)
  pc = PromptController()
  pc.add_handler("ollama", OLLamaTask, Client(host=args.ollama))
  server.register(args.scope, pc)
  server.listen()

if __name__ == "__main__":

  main_cli()
