src/pages/qa_assistant.py

import streamlit as st
from typing import List, Dict
import time

from ..services.ai_analyzer import AIAnalyzer
from ..services.vector_store import VectorStoreService


def show_qa_interface():
    """Display the Q&A assistant interface."""

    if not st.session_state.get("current_document"):
        st.warning("⚠️ No document loaded. Please upload and analyze a document first.")
        if st.button("📄 Go to Upload"):
            st.session_state.page = "📄 Upload"
            st.rerun()
        return

    doc = st.session_state.current_document

    # Header
    st.header("💬 Q&A Assistant")
    st.markdown(f"Ask questions about **{doc.get('filename', 'your document')}**")

    # Initialize chat history
    if "qa_history" not in st.session_state:
        st.session_state.qa_history = []

    # Chat interface
    chat_container = st.container()

    with chat_container:
        # Display chat history
        for i, qa in enumerate(st.session_state.qa_history):
            # User message
            with st.chat_message("user"):
                st.markdown(qa["question"])

            # Assistant response
            with st.chat_message("assistant"):
                st.markdown(qa["answer"])

    # Suggested questions
    st.markdown("### 💡 Suggested Questions")

    doc_type = doc.get("document_type", "other")
    suggested_questions = get_suggested_questions(doc_type)

    col1, col2 = st.columns(2)

    for i, question in enumerate(suggested_questions):
        col = col1 if i % 2 == 0 else col2
        with col:
            if st.button(question, key=f"suggested_{i}", use_container_width=True):
                ask_question(question, doc)

    # Chat input
    st.markdown("### ❓ Ask Your Question")

    with st.form("question_form", clear_on_submit=True):
        user_question = st.text_input(
            "Type your question here...",
            placeholder="e.g., What happens if I terminate this contract early?",
            label_visibility="collapsed",
        )

        submitted = st.form_submit_button("Send", use_container_width=True)

        if submitted and user_question.strip():
            ask_question(user_question, doc)

    # Quick actions
    st.markdown("---")
    col1, col2, col3 = st.columns(3)

    with col1:
        if st.button("📊 Back to Analysis", use_container_width=True):
            st.session_state.page = "📊 Analysis"
            st.rerun()

    with col2:
        if st.button("🗑️ Clear Chat", use_container_width=True):
            st.session_state.qa_history = []
            st.rerun()

    with col3:
        if st.button("📥 Export Chat", use_container_width=True):
            export_chat_history()


def ask_question(question: str, doc: Dict):
    """Process a question and get AI response."""
    try:
        # Show thinking indicator
        with st.spinner("🤔 Thinking..."):
            # Initialize AI analyzer
            ai_analyzer = AIAnalyzer()

            # Get document type
            from ..models.document import DocumentType

            doc_type = DocumentType(doc.get("document_type", "other"))

            # Get answer from AI
            answer = ai_analyzer.answer_question(
                question=question,
                document_text=doc.get("original_text", ""),
                document_type=doc_type,
            )

            # Add to chat history
            st.session_state.qa_history.append(
                {"question": question, "answer": answer, "timestamp": time.time()}
            )

            # Rerun to show the new Q&A
            st.rerun()

    except Exception as e:
        st.error(f"❌ Error processing question: {str(e)}")


def get_suggested_questions(doc_type: str) -> List[str]:
    """Get suggested questions based on document type."""

    questions_by_type = {
        "rental": [
            "What is the monthly rent amount?",
            "What happens if I pay rent late?",
            "How much is the security deposit?",
            "Can I terminate the lease early?",
            "Who is responsible for repairs?",
            "What are the landlord's obligations?",
            "Are pets allowed in the property?",
            "What happens if I damage the property?",
        ],
        "loan": [
            "What is the total amount I will repay?",
            "What is the effective interest rate?",
            "What happens if I miss a payment?",
            "What collateral is required?",
            "Can I repay the loan early?",
            "What are the processing fees?",
            "How is the interest calculated?",
            "What happens in case of default?",
        ],
        "employment": [
            "What is my total compensation package?",
            "How many hours am I expected to work?",
            "Can the company terminate me without notice?",
            "What are the non-compete restrictions?",
            "Am I allowed to work other jobs?",
            "What benefits am I entitled to?",
            "How much notice must I give to resign?",
            "Who owns the intellectual property I create?",
        ],
        "nda": [
            "What information is considered confidential?",
            "How long does the confidentiality last?",
            "What are the penalties for disclosure?",
            "Can I discuss this agreement with others?",
            "What happens after the agreement ends?",
            "Are there any exceptions to confidentiality?",
        ],
        "service": [
            "What services are included in this agreement?",
            "What is the payment schedule?",
            "How can this agreement be terminated?",
            "What are the deliverables and deadlines?",
            "Who is responsible for what costs?",
            "What happens if the work is unsatisfactory?",
        ],
    }

    return questions_by_type.get(
        doc_type,
        [
            "What are the main obligations for each party?",
            "What are the key financial terms?",
            "How can this agreement be terminated?",
            "What are the potential risks for me?",
            "What should I be most careful about?",
            "Are there any unusual or concerning clauses?",
        ],
    )


def export_chat_history():
    """Export the chat history as a text file."""
    if not st.session_state.qa_history:
        st.warning("No chat history to export.")
        return

    doc = st.session_state.current_document

    # Create chat export
    export_text = f"""
LEGA.AI Q&A SESSION EXPORT
{'='*50}

Document: {doc.get('filename', 'Unknown')}
Document Type: {doc.get('document_type', 'Unknown').title()}
Export Date: {time.strftime('%Y-%m-%d %H:%M:%S')}

QUESTIONS & ANSWERS:
{'='*50}

"""

    for i, qa in enumerate(st.session_state.qa_history):
        export_text += f"""
Q{i+1}: {qa['question']}

A{i+1}: {qa['answer']}

{'-'*30}

"""

    export_text += "\nGenerated by Lega.AI - Making legal documents accessible"

    # Clean filename - remove .pdf extension if present
    filename = doc.get("filename", "document")
    if filename.endswith(".pdf"):
        filename = filename[:-4]
    if filename.endswith(".docx"):
        filename = filename[:-5]
    if filename.endswith(".txt"):
        filename = filename[:-4]

    # Offer download
    st.download_button(
        label="📥 Download Chat History",
        data=export_text,
        file_name=f"lega_ai_qa_{filename}.pdf",
        mime="application/pdf",
    )

    st.success("✅ Chat history prepared for download!")