# main.py

import argparse
from pathlib import Path
from rag.vectorstore import FaissStore
from rag.qa import embed_query, answer_question_with_llm
from rag.report import build_report_from_query
from rag.render import render_docx
from rag import config

INDEX_DIR = config.INDEX_DIR


def ask_command(query: str, k: int = 8):
    """
    Q&A mode: retrieve relevant chunks and optionally generate an LLM answer.
    """
    print(f"\n🔍 Question: {query}\n")
    
    store = FaissStore.load(INDEX_DIR)
    qvec = embed_query(query)
    results = store.search(qvec, k=k)

    if not results:
        print("❌ No relevant information found in the documents.")
        return

    print("📄 Retrieved Evidence:\n")
    for idx, r in enumerate(results, 1):
        meta = r["metadata"]
        score = r.get("score", 0.0)
        print(f"[{idx}] {meta.get('doc_name', 'Unknown')} | Page {meta.get('page_number', '?')} | Score: {score:.3f}")
        print(f"    Chunk ID: {meta.get('chunk_id', 'N/A')}")
        print(f"    Text: {meta['text'][:300]}...\n")
        print("-" * 80)

    # Optional: Generate LLM answer
    print("\n🤖 Generating Answer...\n")
    answer = answer_question_with_llm(query, results)
    print(f"✅ Answer:\n{answer}\n")


def report_command(query: str, out_file: str = "outputs/reports/report.docx"):
    """
    Report mode: retrieve relevant chunks, generate structured sections, and render DOCX.
    """
    print(f"\n📝 Generating Report for: {query}\n")
    
    store = FaissStore.load(INDEX_DIR)
    
    # Build report object
    report_obj = build_report_from_query(store, query)
    
    # Render DOCX
    Path(out_file).parent.mkdir(parents=True, exist_ok=True)
    final_path = render_docx(report_obj, out_file)
    
    print(f"\n✅ Report generated successfully at: {final_path}\n")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="RAG Suite CLI - Ask questions or generate reports from indexed PDFs"
    )
    subparsers = parser.add_subparsers(dest="command", required=True)

    # Ask subcommand
    parser_ask = subparsers.add_parser("ask", help="Ask a question (Q&A mode)")
    parser_ask.add_argument("--q", required=True, help="Your question")
    parser_ask.add_argument("--k", type=int, default=8, help="Number of chunks to retrieve (default: 8)")

    # Report subcommand
    parser_report = subparsers.add_parser("report", help="Generate a structured report")
    parser_report.add_argument("--q", required=True, help="Report query/prompt")
    parser_report.add_argument("--out", default="outputs/reports/report.docx", help="Output path for DOCX")

    args = parser.parse_args()

    if args.command == "ask":
        ask_command(args.q, args.k)
    elif args.command == "report":
        report_command(args.q, args.out)