from typing import Dict
from llama_index.core import StorageContext , VectorStoreIndex
from llama_index.core.base.base_retriever import BaseRetriever
from llama_index.core.schema import TextNode , NodeRelationship
from llama_index.embeddings.huggingface_api import HuggingFaceInferenceAPIEmbedding
from chromadb import EphemeralClient
from llama_index.vector_stores.chroma import ChromaVectorStore
from requests import get as reqget
from nest_asyncio import apply
from os import environ
from pickle import load as pickle_load
from bs4 import BeautifulSoup
import json
import flask
class AIBook:
  embed_model:HuggingFaceInferenceAPIEmbedding
  index:VectorStoreIndex
  retriever:BaseRetriever
  def __init__(self,token=environ["hf_api"],srcnum=2):
    self.embed_model = HuggingFaceInferenceAPIEmbedding(model_name="BAAI/bge-large-en-v1.5",token=token)
    self.index = VectorStoreIndex(nodes=pickle_load(open('allbook.book', 'rb')),embed_model=self.embed_model,storage_context=StorageContext.from_defaults(vector_store=ChromaVectorStore(chroma_collection= EphemeralClient().get_or_create_collection("jainbook"))))
    self.retriever = self.index.as_retriever(similarity_top_k=srcnum,vector_store_query_mode="default")
  def changeToken(self,token:str):
    if self.embed_model.token != token:
      self.embed_model = HuggingFaceInferenceAPIEmbedding(model_name="BAAI/bge-large-en-v1.5",token=token)
  def changesrcnum(self,srcnum:int):
    self.retriever = self.index.as_retriever(similarity_top_k=srcnum,vector_store_query_mode="default")
  def retrieve(self,query:str):
    return self.retriever.retrieve(query)
  @classmethod
  def translate_to_eng(cls,text:str)->str:
    return "".join([i[0] for i in reqget(f"https://translate.googleapis.com/translate_a/single?client=gtx&sl=gu&tl=en&dt=t&q={text}").json()[0]])

apply()
app = flask.Flask(__name__)
book = AIBook()
@app.route("/")
def function():
  question = flask.request.args.get("question")
  if question is None:
    return """?question=hi"""
  if question == "":
    return ""
  return flask.redirect(f"/file/images/page-{str(book.retrieve(AIBook.translate_to_eng(question))[0].node.metadata['page']).zfill(3)}.png")

@app.route("/changeToken")
def function2():
  book.changeToken(flask.request.args.get("token"))
  return "Token changed"
@app.route("/changesrcnum")
def function3():
  book.changesrcnum(int(flask.request.args.get("srcnum")))
  return "srcnum changed"
@app.route('/file/images/<filename>')
def serve_image(filename):
    return flask.send_from_directory('images', filename)

if __name__ == '__main__':
    app.run("0.0.0.0",7860)