from flask import Flask, request, jsonify, make_response
from flask_cors import CORS
from langchain_groq import ChatGroq
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain.chains import create_retrieval_chain, create_history_aware_retriever
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory
from dotenv import load_dotenv
import os
import uuid
import traceback

# Load environment variables
load_dotenv()
groq_api_key = os.getenv('GROQ_API_KEY')
os.environ['HF_TOKEN'] = os.getenv('HF_TOKEN')

# Initialize Flask app
app = Flask(__name__)
app.secret_key = os.getenv('FLASK_SECRET_KEY', 'default-secret-key')


CORS(app, supports_credentials=True)



# Initialize LLM and vector store
llm = ChatGroq(model="llama3-70b-8192", groq_api_key=groq_api_key)
embeddings = HuggingFaceEmbeddings(model_name="BAAI/bge-large-en-v1.5")
loaded_db = FAISS.load_local("vectorstore/legal_db", embeddings, allow_dangerous_deserialization=True)
retriever = loaded_db.as_retriever()

# Configure the history-aware retriever
contextualize_q_system_prompt = """Given a chat history and the latest user question,
which might reference context in the chat history, formulate a standalone question
which can be understood without the chat history. Do NOT answer the question,
just reformulate it if needed and otherwise return it as is."""

contextualize_q_prompt = ChatPromptTemplate.from_messages([
    ("system", contextualize_q_system_prompt),
    MessagesPlaceholder("chat_history"),
    ("human", "{input}"),
])

history_aware_retriever = create_history_aware_retriever(
    llm, retriever, contextualize_q_prompt
)

# Configure the question-answering chain
system_prompt = """
You are a highly knowledgeable legal assistant specializing in Pakistani law.
Your responses must use proper legal terminology and cite relevant sections of the law.
If asked about legal actions, suggest options like creating writs or filing appeals.
Be precise and avoid conjecture.
Don't answer any other question other than law related.
Use the following context and chat history to answer the question:
Context: {context}
Chat History: {chat_history}
"""

qa_prompt = ChatPromptTemplate.from_messages([
    ("system", system_prompt),
    MessagesPlaceholder("chat_history"),
    ("human", "{input}"),
])

question_answer_chain = create_stuff_documents_chain(llm, qa_prompt)
rag_chain = create_retrieval_chain(history_aware_retriever, question_answer_chain)

# Session management
store = {}

def get_session_history(session_id: str) -> BaseChatMessageHistory:
    if session_id not in store:
        store[session_id] = ChatMessageHistory()
    return store[session_id]

conversational_rag_chain = RunnableWithMessageHistory(
    rag_chain,
    get_session_history,
    input_messages_key="input",
    history_messages_key="chat_history",
    output_messages_key="answer",
)


@app.route('/test')
def home():
    return 'Backend running!'



@app.route('/api/chat', methods=['POST'])
def chat():
    try:
        data = request.json
        user_input = data.get('input')


        # 1. Get session ID from either:
        #    - Cookies (if browser sent them)
        #    - Headers (fallback for localhost)
        session_id = (
            request.cookies.get('session_id') or 
            request.headers.get('X-Session-ID')
        )
        
        # 1. Get session ID ONLY from cookies (never from JSON body)
        # session_id = request.cookies.get('session_id')
        
        # 2. Create new session ONLY if cookie doesn't exist
        if not session_id:
            session_id = str(uuid.uuid4())

        
        
        # 3. Process the query
        result = conversational_rag_chain.invoke(
            {"input": user_input},
            config={"configurable": {"session_id": session_id}}
        )
        
        
        # 4. Prepare response
        response = jsonify({
            'answer': result['answer'],
            'session_id': session_id  # For debugging only
        })
        
        # 5. Set cookie ONLY if it wasn't already set
        if not request.cookies.get('session_id'):
            response.set_cookie(
                'session_id',
                session_id,
                max_age=30*24*60*60,  # 30 days
                httponly=True,       # Secure against XSS
                secure=True,          # HTTPS only
                samesite='None'       # Allow cross-site usage
                # secure=False,
                # samesite='Lax'
                
            )

        # 6. Critical: Add CORS headers
        response.headers.add('Access-Control-Allow-Origin', 'https://paklaw-chat-advisor.vercel.app/')
        response.headers.add('Access-Control-Allow-Credentials', 'true')
        
        return response
    
    except Exception as e:
        # return jsonify({'error': str(e)}), 500
        # Log the stack trace for debugging
        traceback.print_exc()
        return jsonify({'error': 'Internal Server Error', 'details': str(e)}), 500













        

@app.route('/api/clear_history', methods=['POST'])
def clear_history():
    try:
        session_id = request.cookies.get('session_id')
        if session_id and session_id in store:
            del store[session_id]
            response = jsonify({'status': 'history cleared'})
            response.delete_cookie('session_id')
            return response
        return jsonify({'status': 'no active session'})
    except Exception as e:
        return jsonify({'error': str(e)}), 500