from langchain_core.output_parsers import JsonOutputParser
from langchain_core.prompts import ChatPromptTemplate
from pydantic import BaseModel, Field
from typing import List
from langchain_groq import ChatGroq

import streamlit as st

class RerankResult(BaseModel):
     top_indexes: List[int] = Field(..., description="A releváns dokumentumok indexei")
     message: str = Field(..., description="A felhasználónak megjelenítendő üzenet")
     

class Reranker:
    def __init__(self, model):
        self.llm = model
        self.parser = self.llm.with_structured_output(RerankResult)

    def rerank(self, query: str, docs: List) -> RerankResult:
        context = "\n\n".join(
            [f"[{i}] {doc[0]} (score={doc[1]:.3f})" for i, doc in enumerate(docs)]
        )

        prompt = ChatPromptTemplate.from_template(""" 
Egy intelligens asszisztens vagy, akinek feladata a dokumentumok szűrése azok relevanciája alapján egy kérdés megválaszolásához.
Az alábbi a felhasználó kérdése: "{query}"
És itt vannak a lekért dokumentumok:
{context}

Feladatod:
1. Csak azokat a dokumentumokat válaszd ki, amelyek valóban relevánsak (a többit hagyd figyelmen kívül).
2. Csak a releváns dokumentumok indexeit add vissza (ugyanabban a sorrendben, ahogy a listában szerepelnek).
3. Ha egyik dokumentum sem releváns, akkor a "top_indexes" értéke legyen egy üres lista, a "message" mezőben kérj pontosítást, és a "feedback_required" mező legyen `true`.

⚠️ Nagyon fontos: A válaszod **szigorúan érvényes JSON formátumú legyen**:
- A kulcsok legyenek dupla idézőjelek között (`"`)
- A logikai értékek kisbetűsek legyenek (`true` / `false`)
- Ne használj sortörést a JSON-on kívül

Példa a helyes formátumra:
{{
  "top_indexes": [0, 2],
  "message": ""
  
}}
""")


        # Formatage du prompt
        formatted_prompt = prompt.format_prompt(query=query, context=context)

    # Appel du modèle avec parsing automatique
        result = self.parser.invoke(formatted_prompt.to_messages())
    
        return result