#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from typing import Self, Any
from threading import Thread
from requests import post as Post, Response
from json import loads as json_decode

class AIChat:

    PAYLOAD:dict[str, str|bool] = {
        "model" : "gemma", 
        "stream" : True
    }
    
    def __init__(self:Self) -> None:
        self.__working:bool = True
        self.__thread:Thread = Thread(target=self.__listener)
        self.__thread.start()

    def send(self:Self, message:str) -> str:
        try:

            response:Response

            with Post('http://localhost:11434/api/chat', json = {**self.PAYLOAD, "prompt" : message}) as response:

                line:bytes

                for line in response.iter_lines():
                    if line:

                        chunk:dict[str, Any|None] = json_decode(line)

                        print(chunk.get("response", ""), end = "", flush = True)

                        if chunk.get("done"):
                            break

        except Exception as exception:
            print(f"An error occurred while sending the message: {exception}")

    def __listener(self:Self) -> None:
        while self.__working:
        
            user_input:str = input('> ')

            if user_input in ("close", "exit", "quit"):
                self.__working = False
            else:
                self.send(user_input)

ai_chat = AIChat()