from fasthtml.common import *
from llama_index.core import SimpleDirectoryReader, Document,VectorStoreIndex
from llama_index.core.node_parser import SimpleNodeParser
from llama_index.core.text_splitter import TokenTextSplitter
from llama_index.vector_stores.chroma import ChromaVectorStore
from llama_index.core.storage.storage_context import StorageContext
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.core import SummaryIndex
from llama_index.llms.groq import Groq as GroqLLamaIndex
from chromadb import PersistentClient
from llama_index.core import Settings
from llama_index.embeddings.huggingface_api import (
    HuggingFaceInferenceAPIEmbedding,
)

import chromadb
import os
import threading
import time
from llama_index.core.memory import ChatMemoryBuffer
import json
from llama_index.llms.ollama import Ollama
from llama_index.core.base.response.schema import Response
from groq import Groq
import shutil
pwd = os.getcwd()
def extractConfig(nameModel="SystemData",relPath=os.path.join(pwd,"conf/experiment_config.json"),dataOut="keyantrophics"):
    configPath=os.path.join(os.getcwd(),relPath)
    with open(configPath, 'r', encoding='utf-8') as file:
        config = json.load(file)[nameModel]
    Output= config[dataOut]
    return Output

keygroq=extractConfig(nameModel="SystemData",dataOut="keygroq")
client = Groq(api_key=keygroq)





os.environ["GROQ_API_KEY"] = keygroq


llm_70b = GroqLLamaIndex(model="llama-3.1-70b-versatile")#llm_70b#
llm_localLlamma = llm_70b#Ollama(model="llama3.2")
memory = ChatMemoryBuffer.from_defaults(token_limit=3900)
model_emb="sentence-transformers/multi-qa-mpnet-base-dot-v1"
Settings.llm = llm_localLlamma 



gridlink = Link(rel="stylesheet", href="https://cdnjs.cloudflare.com/ajax/libs/flexboxgrid/6.3.1/flexboxgrid.min.css", type="text/css")

app= FastHTML(hdrs=(picolink, gridlink))



def listUsers():
    with os.scandir("static") as files:
        subdir = [file.name for file in files if  file.is_dir()]
    return subdir

def menuusers(users):
    T=[]
    n=0
    for user in users:
        T.append(Option(user, value=str(user)) )
    return Form(
            H3("Seleccionar usuario",cls="col-xs-3"),
            Select(*T,name="user",cls="col-xs-3"),
            Button("Submit",type="submit",id="buttonMenuuser",cls="col-xs-6"),
            hx_post="/checkInfoSources",hx_swap="innerHTML",hx_target="#files" ,id="menuuser",cls="row middle-xs")

@app.post("/checkInfoSources")
def checkInfoSources(user:str):
    global userdata
    with os.scandir("static/"+user) as files:
        subdir = [Option(file.name,value="static/"+user+"/"+file.name) for file in files if  (file.is_dir() and file.name!="chroma_db") ]
    
    userdata=user
    print("Cambio",userdata)
    return Form(
    H3("Grupos de archivos",cls="col-xs-3"),
    Select(
    *subdir,name="data",cls="col-xs-3"),
    Input(id="name-db", name="collection", placeholder="Enter a collection name",cls="col-xs-4"),
        Button("Submit",type="submit",cls="col-xs-2"), hx_post="/createCollection",hx_swap="innerHTML",hx_target="#status" ,cls="row middle-xs")


def create_or_load_db(path="./chroma_db",collection="init",Nodes=None,modelT=model_emb):
    embed_model = HuggingFaceEmbedding(model_name=modelT)
    #embed_model = HuggingFaceInferenceAPIEmbedding(
    #model_name="BAAI/bge-small-en-v1.5",
    #token="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",  # Optional
    #)
    db = chromadb.PersistentClient(path=path)

    chroma_collection = db.get_or_create_collection(collection)

    vector_store = ChromaVectorStore(chroma_collection=chroma_collection)

    storage_context = StorageContext.from_defaults(vector_store=vector_store)
    if Nodes:

        index = VectorStoreIndex(
            Nodes, storage_context=storage_context, embed_model=embed_model
        )
    else:
        index = VectorStoreIndex.from_vector_store(
            vector_store,
            embed_model=embed_model,
        )
    return index

def post_process_documents(documents):
    processed_documents = []
    n=0
    print(len(documents))
    for doc in documents:
        # 1. Text cleaning
        n+=1
        print(n)
        text = doc.text.lower()  # Convert to lowercase
        # 2. Remove stopwords
        stop_words = set("adssss")
        tokens = text.split(" ")
        filtered_text = ' '.join([word for word in tokens if word.lower() not in stop_words])
        # 3. Custom metadata extraction (example)
        metadata = doc.metadata.copy()
        metadata['word_count'] = len(tokens)
        # 4. Create a new document with processed text and updated metadata
        processed_doc = Document(text=filtered_text, metadata=metadata)
        processed_documents.append(processed_doc)
    node_parser = SimpleNodeParser(chunk_size=360, chunk_overlap=20)
    nodes = node_parser.get_nodes_from_documents(processed_documents)
    return nodes

@app.get("/listmodelactives")
def listmodelactives():
    try:
        print(userdata)
    except:
        print("cambio")
        return Div(id="options",hx_target="this",hx_swap="outerHTML",hx_get="/listmodelactives",hx_trigger="click from:#buttonMenuuser")
    db = chromadb.PersistentClient(path="static/"+userdata+"/chroma_db")
 
    files= db.list_collections()
    collecs = [Option(file.name, value=file.name)for file in files]
    
    return Form(
    Select(
    *collecs,name="data",cls="col-xs-6"),
    Button("Submit",type="submit",cls="col-xs-6"),
    hx_post="/loadCollection",hx_swap="innerHTML",hx_target="#Infomodel")

@app.post("/loadCollection")
def loadCollection(data:str):
    global index
    index=create_or_load_db(path="static/"+userdata+"/chroma_db",collection=data,modelT=model_emb)
    return P("El usuario %s colleccion %s"%(userdata,data))

@app.post("/queryprompt") 
def queryPrompt(question:str):
    #index=load_create_db(collection="my_collection")
    
    query_engine = index.as_query_engine(similarity_top_k=15,vector_store_query_mode="default",response_mode="tree_summarize")
    summary_prompt = (
        "Por favor, genera un resumen completo y detallado del material dado. "
        "Incluye los principales temas, argumentos y conclusiones. "
        "Estructura el resumen de manera coherente y organizada."
    )
    tematic_prompt =(        
        "Por favor, genera un texto donde se menciones la tematica que trata el material dado. "
        "Incluye una tematica general y un indice por temas tratados "
        "Estructura el texto de manera coherente y organizada."

    )
    issues_prompt =(        
        "Por favor, genera un texto donde se menciones tomando en cuenta el material y el contenido de manera detallada que mejoras  podrias realizar al material incluyendo nuevos datos o corrigiendo la informacion proporcionada"
        "Esta mejora hazla pensado paso a paso y de manera muy cuidadosa, tienes que dar ejemplos del materiar o referenciar directamante el texto a mejorar mencioando la causa de la mejora"
        "Estructura el texto de manera coherente y organizada."

    )
    Question_prompt =(        
        "Por favor, genera un texto donde indiques preguntas sobre el material que cuando sean respondidas capturen el punto central y puntos importantes del texto"
        "Estas preguntas hazla pensado paso a paso y de manera muy cuidadosa."
        "Estructura el texto de manera coherente y organizada."

    )
    response = query_engine.query(summary_prompt)
    response2 = query_engine.query(tematic_prompt)
    response3 = query_engine.query(issues_prompt)
    response4 = query_engine.query(Question_prompt)
    Output="<H1>Summary</H1>"+str(response)+"<H1>Tematic</H1>"+str(response2)+"<H1>Issues</H1>"+str(response3)+"<H1>Questions</H1>"+str(response4)

    return Output



@app.post("/chatData")
def questionChat(message:str):
    import logging
    logging.basicConfig(level=logging.INFO)
    contextT=(
            "You are a world-class AI system. You respond to the questions about the context"
            "Here are the relevant documents for the context:\n"
            "{context_str}"
            "\nInstruction: Use the previous chat history, or the context above, carefully examine the given context, to interact and help the user but only about {question}"
            "Never mention the document of reference, talk in first person"        )
    #contextT="Por favor vuelve a repetir el siguiente contenido como tu respuesta:{question}" 
    #'chat_history', 'question'
    query_engine = index.as_query_engine()
    chat_engine = index.as_chat_engine(
        chat_mode="condense_plus_context",
        query_engine=query_engine,
        memory=memory,
        llm=llm_localLlamma,
        context_prompt=contextT,
        similarity_top_k=5,
        verbose=True,
    )
    response = chat_engine.chat(message)
    ContextNodes=""
    for node in response.source_nodes:
        ContextNodes=ContextNodes+node.node.text+"\n"
        #print(f"Texto del nodo: {node.node.text}")
        #print(f"Puntuación de relevancia: {node.score}")
        #print("---")


    
    
    NewPrompt="""The previous response is:
%s
The previous context is:
%s
Evaluate the coherence and accuracy of previous response to respond %s in this evaluation.
Verificate if previous context is related to the previous response, if not, say that you do not have information about that issue 
The format of output is a json with keys 'coherencia', 'exactitud', 'relacion_con_el_contexto' and 'comentario' .
'coherencia', 'exactitud', 'relacion_con_el_contexto' are numeric variables with max value is 10"""%(response,ContextNodes,message)
        


    print(chat_engine.__dict__)

    chat_completion = client.chat.completions.create(
    #
    # Required parameters
    #
    messages=[
        # Set an optional system message. This sets the behavior of the
        # assistant and can be used to provide specific instructions for
        # how it should behave throughout the conversation.
        {
            "role": "system",
            "content": "you are a helpful assistant."
        },
        # Set a user message for the assistant to respond to.
        {
            "role": "user",
            "content": NewPrompt,
        }
    ],

    # The language model which will generate the completion.
    model="llama-3.1-70b-versatile",

    #
    # Optional parameters
    #

    # Controls randomness: lowering results in less random completions.
    # As the temperature approaches zero, the model will become deterministic
    # and repetitive.
    temperature=0.5,

    # The maximum number of tokens to generate. Requests can use up to
    # 32,768 tokens shared between prompt and completion.
    max_tokens=1024,

    # Controls diversity via nucleus sampling: 0.5 means half of all
    # likelihood-weighted options are considered.
    top_p=1,

    # A stop sequence is a predefined or user-specified text string that
    # signals an AI to stop generating content, ensuring its responses
    # remain focused and concise. Examples include punctuation marks and
    # markers like "[end]".
    stop=None,

    # If set, partial message deltas will be sent.
    stream=False,
    )


    return P(message),P(response),P(chat_completion.choices[0].message.content)

@app.get("/SummarySources")
def SummarySources():
    with os.scandir("static/"+userdata) as files:
        subdir = [Option(file.name,value="static/"+userdata+"/"+file.name) for file in files if  file.is_file()]
    return Form("Este es muy caro para documentos grandes y tarda mucho",
    Select(
    *subdir,name="data"),
    Input( name="query", placeholder="Enter a query"),
        Button("Submit",type="submit"), hx_post="/SummaryMake",hx_swap="innerHTML",hx_target="#summaryR" )

@app.post("/SummaryMake")
def SummaryMake(data:str,query:str):
    print(data,query)
    docs = SimpleDirectoryReader(
        input_dir=[data]
    ).load_data()
    print("p1")
    summary_index = SummaryIndex.from_documents(docs)
    print("p2")
    summary_engine = summary_index.as_query_engine()
    print("p3")
    response = summary_engine.query(
    query
    )
    print("p4")
    return P(response)



@app.post("/createCollection")
def createCollection(data:str,collection:str):
    print("Reading")
    docs = SimpleDirectoryReader(
        input_dir=data
    ).load_data()
    print("Process Documents")
    Nodes=post_process_documents(docs)
    print("create DB")
    class MyThread(threading.Thread):
        def run(self):
             print("Hilo")
             create_or_load_db(path="static/"+data.split("/")[1]+"/chroma_db",collection=collection,Nodes=Nodes,modelT=model_emb)

    # create and start the thread
    global t
    t = MyThread()
    t.start()
    global t_time
    t_time=time.time()
    return Div("Iniciando carga de datos")
 
@app.get("/is_busy")
def is_busy():
    try:
        Busy= t.is_alive()  
    except:
        Busy=False
    if not Busy:
        return H2("Estado: Disponible para carga de datos")
    else: 
        return H2("Esta ocupados desde hace %s , este es un proceso largo"%(str(time.time()-t_time)))



@app.get("/")
def home():
    
    page = Title('Super tutor'),Main(
        Div(H1('Super tutor'), 
                menuusers(listUsers()),
                #A('A link', href='https://example.com'), 
                #Img(src="https://placehold.co/200"),
                Div("Archivos",id="files"),
                
                Div(H2("Estado:Disponible para carga"),id="status",hx_target="this",hx_swap="innerHTML",hx_get="/is_busy",hx_trigger="every 60000ms"),
                Div(
                    Div(Div(id="options",hx_target="this",hx_swap="outerHTML",hx_get="/listmodelactives",hx_trigger="click from:#buttonMenuuser delay:3s"),cls="col-xs-12"),
                    Div(Div(id="Infomodel"),cls="col-xs-12"),
                    #Div("Resumen",Div(id="summary",hx_target="this",hx_swap="outerHTML",hx_get="/SummarySources",hx_trigger="click from:#buttonMenuuser"),Div(id="summaryR")),
                    Div(
                        Div(
                            Form(
                                Input(id="question", name="message", placeholder="Enter a message"),
                                Button("Submit",type="submit"), hx_post="/chatData",hx_swap="afterend",hx_target="#questionR" ),
                            Div(id="questionR")
                        ,id="questions"),
                    cls="col-xs-6"),
                    Div(
                        Div(
                            Form(
                                Input(id="query", name="question", placeholder="Enter a query"),
                                Button("Submit",type="submit"), hx_post="/queryprompt",hx_swap="innerHTML",hx_target="#queryR" ),
                            Div(id="queryR"),
                            id="query"),
                        id="chatbot",cls="col-xs-6"),
                    cls="row", style="color: #fff;")
                ))
    return page

# @app.post("/upload")
# def upload(data: UploadFile = File(...),user : str = Form(...), dir: str = Form(...)):
#     filename="static/"+user+dir+data.filename

#     with open(f"{filename}", "wb") as buffer:
#         shutil.copyfileobj(data.file, buffer)

# app.mount("/static", StaticFiles(directory="static"), name="static")




serve()