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

from typing import Self, Any, Sequence, Optional, Callable
from threading import Thread
from requests import post as Post, Response
from json import loads as json_decode
from re import compile as re_compile, Pattern as REPattern, IGNORECASE as RE_IGNORECASE
from websockets.sync import server as WebSocketServer
from Utils.Utils import Utils
from Utils.Check import Check
from Drivers.FilesDriver import FilesDriver

class PayloadModel:
    def __init__(self:Self) -> None:
        self.model:str = "gemma"
        self.stream:bool = True

class AIChat:

    DEFAULT_SETTINGS:dict[str, Any|None] = {
        "autostart" : True, 
        "ai_model" : "gemma", 
        "ai_stream" : True
    }
    
    def __init__(self:Self, inputs:Optional[dict[str, Any|None]|Sequence[Any|None]] = None) -> None:

        self.__inputs:dict[str, Any|None] = self.get_dictionary(inputs, overwrite = True)
        self.__working:bool = False
        self.__started:bool = False
        self.__thread:Thread|None = None
        self.__payload:PayloadModel = PayloadModel()

        self.files:FilesDriver = FilesDriver(self)

    def start(self:Self, callback:Optional[Callable[[bool], bool]] = None) -> bool:
        
        end:Callable[[bool], bool] = lambda ok: callback(ok) if callable(callback) else ok

        if self.__started:
            return end(False)
        self.__started = True

        self.__working = True
        self.__payload.model = self.get("ai_model", self.__payload.model)
        self.__thread = Thread(target = self.__listener)
        self.__thread.start()

        return end(True)

    def close(self:Self, callback:Optional[Callable[[bool], bool]] = None) -> bool:
        
        end:Callable[[bool], bool] = lambda ok: callback(ok) if callable(callback) else ok

        if not self.__started:
            return end(False)
        self.__started = False

        self.__working = False

        return end(True)
    
    def load_json(self:Self, data:Any|None, only_dictionary:bool = True) -> list[list[Any|None]|dict[str, Any|None]]:

        items:list[list[Any|None]|dict[str, Any|None]] = []

        if Check.is_string(data):
            items.extends(self.files.load_json(data))
        elif Utils.is_array(data):
            if only_dictionary:

                item:Any|None

                for item in data:
                    items.extend(self.load_json(item, only_dictionary))

            else:
                items.append(data)
        elif Utils.is_dictionary(data):
            items.append(data)

        return items

    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()