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 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 nodes_to_guj(cls,nodes:list[TextNode])->list[str]:
    return [node.node.relationships[NodeRelationship("1")].metadata["maintext"] for node in nodes]
  @classmethod
  def nodes_to_eng(cls,nodes:list[TextNode])->list[str]:
    return [node.node.text for node in nodes]
  @classmethod
  def nodes_to_page_with_bookname(cls,nodes:list[TextNode])->list[Dict]:
    return [{"page":int(node.node.relationships[NodeRelationship("1")].metadata["page"]),"bookname":node.node.relationships[NodeRelationship("1")].metadata["book"]} for node in nodes]
  @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 """Please provide a question as 'https://shethjenil-apiofbookai.hf.space?question=પ્રભુ છે કે નહિ'<br>for getting image of question https://shethjenil-apiofbookai.hf.space/question_to_img?question=પ્રભુ છે કે નહિ<br>for get image by book and page https://shethjenil-apiofbookai.hf.space/getbookimage/023657/99<br>for knowing full data with page etc https://shethjenil-apiofbookai.hf.space/fulldetails?question=પ્રભુ છે કે નહિ<br>Change Token as https://shethjenil-apiofbookai.hf.space/changeToken?token=hf_<b>rgPNhjnXpLSodIphwjmRvPbvrovNYnQavj</b><br>Change srcnum as https://shethjenil-apiofbookai.hf.space/changesrcnum?srcnum=2<br><br><form action="/insert_node" method="post" enctype="multipart/form-data"><input type="file" name="file" accept=".book"><input type="submit" value="Upload Book File"></form>"""
  if question == "":
    return ""
  return "\n\n".join(AIBook.nodes_to_guj(book.retrieve(AIBook.translate_to_eng(question))))

@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("/insert_nodes",methods=["POST"])
def function4():
  if 'file' not in flask.request.files:
    return "Error"
  file = flask.request.files['file']
  if file.filename.endswith(".book"):
    book.index.insert(pickle_load(file))
    return "inserted" + file.filename
  else:
    return "not inserted because file is not a book"
@app.route("/fulldetails")
def function5():
  question = flask.request.args.get("question")
  if question:
    return flask.jsonify([{"text":i.get_text(),"score":i.get_score(),"metadata":i.node.relationships[NodeRelationship("1")].metadata} for i in book.retrieve(AIBook.translate_to_eng(question))])
  else:
    return "please provide a question parameter https://shethjenil-apiofbookai.hf.space/fulldetails?question=પ્રભુ છે કે નહિ"
@app.route("/getbookimage/<bookid>/<page>")
def function6(bookid,page):
#   response = reqget(BeautifulSoup(reqget(f"https://jainqq.org/explore/{bookid}/{page}").content, "html.parser").find("img",class_="img-fluid").get("src"))
#   return flask.send_file(BytesIO(response.content), mimetype=response.headers['Content-Type'])
  return flask.redirect(BeautifulSoup(reqget(f"https://jainqq.org/explore/{bookid}/{page}").content, "html.parser").find("img",class_="img-fluid").get("src"))
@app.route("/question_to_img")
def function7():
  question = flask.request.args.get("question")
  if question:
    meta = book.retrieve(AIBook.translate_to_eng(question))[0].node.relationships[NodeRelationship("1")].metadata
    flask.redirect(BeautifulSoup(reqget(f"https://jainqq.org/explore/{meta['bookid']}/{meta['page']}").content, "html.parser").find("img",class_="img-fluid").get("src"))

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