from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.chains.retrieval import create_retrieval_chain
from langchain_groq import ChatGroq
from langchain_core.prompts import ChatPromptTemplate
import os
from langchain_core.prompts import PromptTemplate
import logging
from embeddings.embedder import get_vectorstore


logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelName)s - %(message)s'
)
logger = logging.getLogger(__name__)


rag_prompt_template = '''
You are a legal assistant trained to answer questions using legal documents.

If the answer cannot be determined from the available legal text, you must answer without include this phrase:
  > According to Article <article>, Chapter <chapter>, <title> 
For example:
   > The answer cannot be determined from the available legal text.

Otherwise, you must STRICTLY follow this 3-step structure:

1. **Begin your response ONLY with this format** (fill in the values from metadata):
   > According to Article <article>, Chapter <chapter>, <title>:


2. **Then, extract and summarize the most relevant point** from the provided context.

3. **Finally, answer the user’s question clearly and formally**, referring only to the point above.

---
💥 IMPORTANT RULES:
- Do NOT repeat or paraphrase the law reference later.
- Do NOT invent any legal information — use ONLY the provided context and metadata.
- Do NOT add phrase like 'The most relevant point is that' when mentioning the context.
- If the context is insufficient, nrever include the legal citation and respond with:
  > "The answer cannot be determined from the available legal text."
---

**User Question:**  
{input}

**Retrieved Legal Context:**  
{context}
'''

document_prompt = PromptTemplate.from_template(
    '''
    Article: {article}
    Chapter: {chapter}
    Title: {title}
    Content: {page_content}
    '''
)

prompt = ChatPromptTemplate.from_messages([
    ('system', rag_prompt_template),
    ('user', "Context:\n{context}\n\nQuestion:\n{input}\n\nAnswer:")
])

def get_rag_chain(model_name='llama-3.3-70b-versatile', k=3):
    db = get_vectorstore()

    groq_api_key = os.getenv('GROQ_API_KEY')

    llm = ChatGroq(
    model_name=model_name,
    temperature=0,
    max_tokens=10000,
    api_key=groq_api_key,
)

    retriever = db.as_retriever(
        search_type = 'similarity',
        search_kwargs = {'k': k}
    )

    combine_doc_chain = create_stuff_documents_chain(
        prompt=prompt,
        llm=llm,
        document_prompt=document_prompt,
    )
    return create_retrieval_chain(retriever, combine_doc_chain)