from dotenv import load_dotenv
from together import Together
from sentence_transformers import SentenceTransformer
from sentence_transformers import SentenceTransformer
import vecs
import gradio as gr
import os 

load_dotenv()

user = os.getenv("user")
password = os.getenv("password")
host = os.getenv("host")
port = os.getenv("port")
db_name = "postgres"
DB_CONNECTION = f"postgresql://{user}:{password}@{host}:{port}/{db_name}"
vx = vecs.create_client(DB_CONNECTION)
model = SentenceTransformer('Snowflake/snowflake-arctic-embed-xs')
client = Together(api_key=os.getenv('TOGETHER_API_KEY'))

def query_db(query, limit = 5, filters = {}, measure = "cosine_distance", include_value = False, include_metadata=False, table = "2023"):
  query_embeds = vx.get_or_create_collection(name= table, dimension=384)
  ans = query_embeds.query(
      data=query,
      limit=limit,
      filters=filters,
      measure=measure,
      include_value=include_value,
      include_metadata=include_metadata,
  )
  return ans  

def construct_result(ans):
  results = ""
  for i in range(0, len(ans)):
    a, b = ans[i][2].get("sentencia"), ans[i][2].get("fragmento")
    results += (f"En la sentencia {a}, se dijo {b}\n")
  return results

def sort_by_score(item):
  return item[1]

def referencias(results):
  references = 'Sentencias encontradas: \n'
  enlistadas = []
  for item in results: 
    if item[2].get('sentencia') in enlistadas: 
      pass
    else: 
      references += item[2].get('sentencia')+ ' '
      enlistadas.append(item[2].get('sentencia'))
  return references

def live_inference(messages, max_new_tokens = 512):
    response = client.chat.completions.create(
    model="meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo",
    messages = messages,
    max_tokens = max_new_tokens
    )
    return response.choices[0].message.content

def inference(hechos, investigacion, teoria):
  encoded_prompt1 = model.encode(investigacion)
  years = range(2019, 2025)
  results = []
  for year in years:
    results.extend(query_db(encoded_prompt1, include_metadata = True, table = str(year), include_value=True, limit = 5))
  results.sort(key=sort_by_score, reverse=True)
  researchAI=[
        {"role": "system", "content": f"""
                Eres Ticio, un asistente de investigación jurídica. Tu deber es organizar el contenido de las sentencias de la jurisprudencia de acuerdo 
                a las necesidades del usuario. Debes responder solo en español. Debes responder solo en base a la información del contexto a continuación.
                Siempre debes mencionar la fuente en tu escrito, debe tener un estilo formal y juridico.
                Contexto: 
                {construct_result(results)}
                """
        },
        {"role": "user", "content": investigacion},
    ]
  research = live_inference(researchAI, max_new_tokens=1024)
  factsAI =[
        {"role": "system", "content": f"""
                Eres Ticio, un asistente de redacción jurídica. Debes ayudar a organizar el texto de los hechos. 
                el texto debe verse como una lista numerada de parrafos con la siguiente estructura: 
                numero. fecha. descripción del evento. 
                Debes usar lenguaje juridico y profesional para explicar el evento y explicarlos de la forma que mejor beneficie 
                la teoría del caso del usuario.
                Teoría del caso:
                {teoria}
                """
        },
        {"role": "user", "content": hechos},
    ]
  facts = live_inference(factsAI, max_new_tokens=1024)
  argsAI = [
        {"role": "system", "content": f"""
                Eres Ticio, un asistente de redacción jurídica. Tu tarea es redactar los fundamentos de derecho de la tutela del usuario.
                Para redactar unos fundamentos de derecho debes escribir varios argumentos en uno o dos parrafos de la siguiente forma:
                Fuente juridica (artiuclo de un decreto, ley o resolución).
                Estandar jurídico (debes mencionar una sentencia judicial que interprete la fuente juridica y explique como se debe aplicar). 
                Aplicación al caso (debes explicar en varios parrafos el analisis detallados donde menciones las fuentes y estandares porque los hechos del caso se ajustan al estandar que mencionaste antes).

                Redacta la aplicación del caso en parrafos continuos sin bullet points
                Debes utilizar las fuentes y estandares juridicos que aparezcan en la investigación que te proporciona el usuario
                El argumento debe ayudar a defender la postura de la teoría del caso
                Teoría del caso:
                {teoria}

                Hechos:
                {facts}

                recuerda redactar de forma que expliques en lenguaje juridico y técnico el analisis de cada uno de los argumentos 
                """
        },
        {"role": "user", "content": research},
    ]
  args = live_inference(argsAI, max_new_tokens=1536)
  return "**Hechos**: \n\n" + facts + "\n\n **Fundamentos**: \n\n" + args
  
theme = gr.themes.Base(
    primary_hue="red",
    secondary_hue="red",
    neutral_hue="neutral",
).set(
    button_primary_background_fill='#910A0A',
    button_primary_border_color='*primary_300',
    button_primary_text_color='*primary_50'
)


with gr.Blocks(theme=theme) as demo:
      hechos = gr.Textbox(label = "Hechos", lines = 15, show_label = True, show_copy_button= True)
      tema = gr.Textbox(label = "Investigación", lines = 1, show_label = True, show_copy_button= True)
      teoria = gr.Textbox(label = "Teoría del caso", lines = 5, show_label = True, show_copy_button= True)
      btn = gr.Button("Generar")
      output = gr.Markdown(label = "respuesta", show_label = True, show_copy_button= True)
      btn.click(fn=inference, inputs=[hechos, tema, teoria], outputs=output, api_name=False)

if __name__ == "__main__":
    demo.queue()
    demo.launch(show_api=False, debug= True)