import openai
import llama_index
import cryptography
import cryptography.fernet
import huggingface_hub
import huggingface_hub.hf_api
import json
import os
import gradio
class HFace_Pluto(object):
  # initialize the object
  def __init__(self, name="Pluto",*args, **kwargs):
    super(HFace_Pluto, self).__init__(*args, **kwargs)
    self.author = "Duc Haba, Girish"
    self.name = name
    self._ph()
    self._pp("Hello from class", str(self.__class__) + " Class: " + str(self.__class__.__name__))
    self._pp("Code name", self.name)
    self._pp("Author is", self.author)
    self._ph()
    #
    # define class var for stable division
    self._device = 'cuda'
    self._steps = [3,8,21,55,89,144]
    self._guidances = [1.1,3.0,5.0,8.0,13.0,21.0]
    self._xkeyfile = '.xoxo'
    self._models = []
    self._seed = 667 # sum of walnut in ascii (or Angle 667)
    self._width = 512
    self._height = 512
    self._step = 50
    self._guidances = 7.5
    self._llama_query_engine = None
    self._llama_index_doc = None
    #self._generator = torch.Generator(device='cuda')
    self.pipes = []
    self.prompts = []
    self.images = []
    self.seeds = []
    self.fname_id = 0
    self.dname_img = "img_colab/"
    self._huggingface_key="gAAAAABkgtmOIjpnjwXFWmgh1j2et2kMjHUze-ym6h3BieAp34Sqkqv3EVYvRinETvpw-kXu7RSRl5_9FqrYe-7unfakMvMkU8nHrfB3hBSC76ZTXwkVSzlN0RfBNs9NL8BGjaSJ8mz8"
    self._gpt_key="'gAAAAABkgtoTOLPegnxNIAfBfAda17h5HIHTS_65bobO3SdDlJam07AHGrcolvk9c6IWNJtTTxaCb8_JtWnLz0Y5h9doyfL-nJZggeQ6kLtaD4XwZYcG-AtYNNGCnJzVt9AaysPDnu-KWVhnJSe-DyH0oOO33doE0g=='"
    self._fkey="=cvsOPRcWD6JONmdr4Sh6-PqF6nT1InYh965mI8f_sef"
    self._color_primary = '#2780e3' #blue
    self._color_secondary = '#373a3c' #dark gray
    self._color_success = '#3fb618' #green
    self._color_info = '#9954bb' #purple
    self._color_warning = '#ff7518' #orange
    self._color_danger = '#ff0039' #red
    self._color_mid_gray = '#495057'
    return
  #
  # pretty print output name-value line
  def _pp(self, a, b,is_print=True):
    # print("%34s : %s" % (str(a), str(b)))
    x = f'{"%34s" % str(a)} : {str(b)}'
    y = None
    if (is_print):
      print(x)
    else:
      y = x
    return y
  #
  # pretty print the header or footer lines
  def _ph(self,is_print=True):
    x = f'{"-"*34} : {"-"*34}'
    y = None
    if (is_print):
      print(x)
    else:
      y = x
    return y
  #
  def push_hface_files(self,
    hf_names,
    hf_space="girishlkiran/ct",
    local_dir="/content/"):
    f = str(hf_names) + " is not iteratable, type: " + str(type(hf_names))
    try:
      for f in hf_names:
        lo = local_dir + f
        huggingface_hub.upload_file(
          path_or_fileobj=lo,
          path_in_repo=f,
          repo_id=hf_space,
          repo_type=huggingface_hub.REPO_TYPE_SPACE)
    except Exception as e:
      self._pp("*Error", e)
    return
  #
  def push_hface_folder(self, hf_folder, hf_space_id, hf_dest_folder=None):
    api = huggingface_hub.HfApi()
    api.upload_folder(folder_path=hf_folder,
      repo_id=hf_space_id,
      path_in_repo=hf_dest_folder,
      repo_type="space")
    return
  #
  def write_file(self,fname, txt):
    f = open(fname, "w")
    f.writelines("\n".join(txt))
    f.close()
    return
  #
  def _fetch_crypt(self,is_generate=False):
    s=self._fkey[::-1]
    if (is_generate):
      s=open(self._xkeyfile, "rb").read()
    return s
  #
  def _gen_key(self):
    key = cryptography.fernet.Fernet.generate_key()
    with open(self._xkeyfile, "wb") as key_file:
        key_file.write(key)
    return
  #
  def _decrypt_it(self, x):
    y = self._fetch_crypt()
    f = cryptography.fernet.Fernet(y)
    m = f.decrypt(x)
    return m.decode()
  #
  def _encrypt_it(self, x):
    key = self._fetch_crypt()
    p = x.encode()
    f = cryptography.fernet.Fernet(key)
    y = f.encrypt(p)
    return y
  #
  def _login_hface(self):
    huggingface_hub.login(self._decrypt_it(self._huggingface_key),
      add_to_git_credential=True) # non-blocking login
    self._ph()
    return
  #
  def _setup_openai(self,key=None):
    if (key is None):
      key = self._decrypt_it(self._gpt_key)
    #
    openai.api_key = key
    os.environ["OPENAI_API_KEY"] = key
    return
  #
  def _fetch_index_files(self,llama_ix):
    res = []
    x = llama_ix.ref_doc_info
    for val in x.values():
      jdata = json.loads(val.to_json())
      fname = jdata['extra_info']['file_name']
      res.append(fname)
    # remove dublication name
    res = list(set(res))
    return res
# add module/method
#
import functools
def add_method(cls):
  def decorator(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
      return func(*args, **kwargs)
    setattr(cls, func.__name__, wrapper)
    return func # returning func means func can still be used normally
  return decorator
#
monty = HFace_Pluto("Monty")
monty._login_hface()
monty._ph()
monty._setup_openai()

@add_method(HFace_Pluto)
def load_llama_index(self,vindex='vector_index',vpath='./index_storage'):
  try:
    storage_context = llama_index.StorageContext.from_defaults(persist_dir=vpath)
    # load index
    self._llama_index_doc = llama_index.load_index_from_storage(storage_context, index_id=vindex)
    print(f'Index doc are: {self._fetch_index_files(self._llama_index_doc)}')
  except Exception as e:
    print('**Error: can not load index, check the index_storage directory or the GPT auth token')
    print('If do not have index tokens then run the .gen_llama_index() function')
    print(f'Exception: {e}')
  return
monty.load_llama_index()

@add_method(HFace_Pluto)
def ask_me(self, p, ll_sign_in_member='Girish' , ll_engine='Humana'):
  self._llama_query_engine = self._llama_index_doc.as_query_engine()
  ll_engine = self._llama_query_engine

  px = f'My name is {ll_sign_in_member}, and I want answer to the following: {p}.'
  print("##### " + px)
  
  resp = ll_engine.query(px)
  return resp

in_box = [gradio.Textbox(lines=1, label="Your Humana request", placeholder="Your Humana request...see example if you need help.")
  ,gradio.Radio(["Girish"], label="Login Member", value='Girish', info="Who had login?")
  ,gradio.Radio(["Humana"], label="Login Member", value='Humana', info="Fine-Tune LLM for:")
  ]
out_box = [gradio.Textbox(label="Humana response:")]
#

title = "Humana and C&T Fine-tune LLM model"
desc = '*Note: This model is fine-tuned by YML using GPT3.5 as the base LLM.'
arti = '<li><i>**Note: You can add more documentation. The more documentation the model has the smarter it will be.</i></li></ul>'
exp = [
  ['Tell me the Humana Gold Plus plan.'],
  ['Please write a summary in bullet point of the Humana Gold Plus SNP-DE H0028-015 (HMO-POS D-SNP) Annual Notice of Changes for 2023.'],
  ['Write a newsletter introducing Humana Gold Plus plan, and target it to senior citizen demographic.'],
  ['Please write a summary about the Humana and Longevity Health Partner so that a teenage can understand.'],
  ['Tell me about the state agency contact information in bullet point.'],
  ['Please tell me more about the Humana Offer Free Counseling about Medicare and Medicaid'],
  ['Write four engaging tweets about the Humana Gold Plus plan.'],
  ['Is Humana is same as human?']
 ]
flag_opt = [': Good', ': Bad']
flag_dir = './user_feed_back'

gradio.Interface(fn=monty.ask_me,
  inputs=in_box,
  outputs=out_box,
  examples=exp,
  title=title,
  description=desc,
  article=arti).launch(debug=True)