from backend import get_handler
import uuid
import os
import json
import time
from datetime import datetime
from app_utils import *

class State:
  """
    Manages the state of a chatbot including its model configuration, chat history, and test data.
    
    Attributes:
        current_model (str): The currently active model identifier
        current_temperature (float): The temperature setting for model inference
        current_category (str): The current category of conversation
        test_entry (dict): Contains test configuration and data
        inference_data (dict): Stores inference results and messages
        handler: The model handler instance
        history (list): List of conversation messages
        name (str): Identifier for this bot instance, defaults to "bot"
        database: Database connection for storing chat history
        example_settings (dict): Predefined example configurations (see app_utils.py)
    """
  def __init__(self, init_model, init_temperature, database, example_settings, name="bot"):
      self.current_model = init_model
      self.current_temperature = init_temperature
      self.current_category = None
      self.test_entry = initialize_empty_test_entry()
      self.inference_data = {"message": []}
      self.handler = get_handler(self.current_model, self.current_temperature)
      self.history = self.get_initial_state()
      self.name = name
      self.database = database
      self.example_settings = example_settings
      
  def initialize(self, category):
      self.category = category
      self.update_category_and_load_config(category)
    
  def get_initial_state(self):
      return list(INITIAL_CHAT_HISTORY)
      
  def restart_chat(self, model, temperature):
      self.test_entry["id"] = str(uuid.uuid4()) 
      self.handler = get_handler(model, temperature)
      self.history = self.get_initial_state()
      return self.history
    
  def update_handler(self, model, temperature):
      self.current_model = model
      self.current_temperature = temperature
      self.handler = get_handler(model, temperature)
      print(f"Update handler for {self.name}: ", model, temperature)
      self.history = self.restart_chat(model, temperature)
      return model, self.history
    
  def save(self):
      if len(self.history) > 2:
          document = {"time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                      "model": self.current_model,
                      "category": self.current_category,
                      "temperature": self.current_temperature,
                      "history": self.history}
          self.database.insert_one(document)
          
  def restart_chat_and_save(self):
      self.save()
      return self.restart_chat(self.current_model, self.current_temperature)
      
  def update_category_and_load_config(self, category):
      self.current_category = category
      self.test_entry["initial_config"] = {category: {}}
      self.test_entry["involved_classes"] = [category]
      config_path = os.path.join("config", f"{MAPPINGS[category]}.json")
      self.load_config(config_path)
      return category
      
  def load_config(self, config_path):
      if os.path.exists(config_path):
          with open(config_path, 'r') as config_file:
              data = json.load(config_file)
              self.test_entry["function"] = data.copy()
              
  def load_example_and_update(self, example):
      self.save()
      model, temp, category, message = self.load_example(example)
      self.update_category_and_load_config(category)
      return model, temp, category, message
      
  def load_example(self, example):
      return self.example_settings[example]
    
  def response(self):
      for item in self.handler.inference(self.test_entry):
          if item[0] == "regular":
              responses_results = equalize_and_zip(item[1], item[2])
              for (model_res, exec_res) in responses_results:
                  if model_res is not None:
                      response = model_res
                      self.history.append({"role": "assistant", "content": "<b>Model Response🤖: </b><br>"})
                      for character in response:
                          self.history[-1]["content"] += character
                          time.sleep(0.01)
                          yield self.history
                  if exec_res is not None:
                      response = exec_res
                      self.history[-1]["content"] += "<br><br><b>Model Execution💻: </b><br>"
                      yield self.history
                      for character in response:
                          self.history[-1]["content"] += character
                          time.sleep(0.01)
                          yield self.history
          elif item[0] == 'summary':
              response = item[1]
              if response is not None:
                  self.history.append({"role": "assistant", "content": "<b>Summary✅: </b><br>"})
                  for character in response:
                      self.history[-1]["content"] += character
                      time.sleep(0.01)
                      yield self.history
          elif item[0] == "final":
              self.inference_data = item[2] 
          time.sleep(0.05)