main.py

#!/usr/bin/env python3
"""
RAG System CLI - Main entry point.

Supports document processing, interactive queries, batch processing, and summaries.

Usage:
    python main.py --docs data/ --interactive    # Process docs + interactive mode
    python main.py --docs data/ --query "..."    # Single query
    python main.py --batch queries.txt           # Batch queries from file
    python main.py --docs data/ --summary        # Generate document summary
"""

import argparse
import sys
from pathlib import Path

from dotenv import load_dotenv

load_dotenv()

from config.config import get_default_config
from src.pipeline.advanced_rag_pipeline import create_rag_pipeline
from src.utils.helpers import validate_api_keys
from src.utils.logging_config import get_logger, setup_logging


def setup_environment():
    """Initialize logging and validate required API keys.

    Returns:
        Tuple of (success: bool, logger)
    """
    logger = setup_logging(log_level="INFO", console_output=True)

    missing = [k for k, v in validate_api_keys(["GOOGLE_API_KEY"]).items() if not v]
    if missing:
        logger.error(f"Missing API keys: {', '.join(missing)}")
        print("\n❌ Configurazione incompleta!")
        print("Configura GOOGLE_API_KEY nel file .env (vedi .env.example)")
        return False, logger

    logger.info("✅ Ambiente configurato correttamente")
    return True, logger


def process_documents(pipeline, document_paths, force_reprocess=False):
    """Process documents through the pipeline.

    Returns:
        True if successful, False otherwise
    """
    try:
        docs = pipeline.process_documents(
            document_paths, force_reprocess=force_reprocess
        )
        print(f"✅ Processati {len(docs)} documenti")
        return True
    except Exception as e:
        print(f"❌ Errore nel processing: {e}")
        return False


def interactive_mode(pipeline):
    """Modalità interattiva per query"""
    print("\n" + "=" * 60)
    print("💬 MODALITÀ INTERATTIVA")
    print("=" * 60)
    print("Comandi disponibili:")
    print("  - Scrivi una domanda per ottenere una risposta")
    print("  - 'summary' per un riassunto dei documenti")
    print("  - 'stats' per statistiche del sistema")
    print("  - 'help' per questo messaggio")
    print("  - 'quit' per uscire")
    print("-" * 60)

    # Logger per scrivere anche su file nella cartella logs/
    rag_logger = get_logger("rag_system")

    while True:
        try:
            user_input = input("\n🔍 > ").strip()

            if user_input.lower() in ["quit", "exit", "q"]:
                break
            elif user_input.lower() == "help":
                print("Comandi: query normale, 'summary', 'stats', 'quit'")
                continue
            elif user_input.lower() == "summary":
                print("📄 Generando riassunto...")
                try:
                    summary = pipeline.get_document_summary("comprehensive")
                    print(f"\n📋 Riassunto:\n{summary.answer}")
                except Exception as e:
                    print(f"❌ Errore nel riassunto: {e}")
                continue
            elif user_input.lower() == "stats":
                if pipeline.state:
                    print(f"\n📊 Statistiche:")
                    print(
                        f"  - Documenti processati: {len(pipeline.state.processed_documents)}"
                    )
                    print(
                        f"  - Chunks semantici: {len(pipeline.state.semantic_chunks)}"
                    )
                    print(
                        f"  - Chunks arricchiti: {len(pipeline.state.enriched_chunks)}"
                    )
                    total_questions = sum(
                        len(chunk.hypothetical_questions)
                        for chunk in pipeline.state.enriched_chunks
                    )
                    print(f"  - Domande ipotetiche: {total_questions}")
                else:
                    print("❌ Nessun dato processato")
                continue
            elif not user_input:
                continue

            # Esegui query normale
            print("🤔 Elaborando...")

            # Oggetto RagResponse
            response = pipeline.query(user_input)

            print(f"\n🎯 Risposta (Confidence: {response.confidence:.2f}):")
            print(response.answer)

            if response.sources:
                print(f"\n📚 Fonti principali:")
                # Mappa chunk_id -> risultato di retrieval per accedere al contenuto
                results_by_id = {
                    r.chunk_id: r for r in getattr(response, "retrieval_results", [])
                }
                for i, source in enumerate(
                    response.sources[: get_default_config().fusion_retrieval.top_k], 1
                ):
                    print(f"  {i}. {source}")
                    # Logga anche su file
                    rag_logger.info(f"Fonte {i}: {source}")
                    chunk = results_by_id.get(source)
                    if chunk and getattr(chunk, "content", None):
                        # Mostra un estratto del contenuto per evitare output troppo lungo
                        snippet = chunk.content.strip().replace("\n", " ")
                        score_info = (
                            f" (score: {getattr(chunk, 'score', 0):.3f})"
                            if hasattr(chunk, "score")
                            else ""
                        )
                        print(f"     └─ Estratto{score_info}: {snippet}")
                        # E scrivilo anche nei log (snippet completo potrebbe essere lungo, quindi lo mettiamo a DEBUG)
                        rag_logger.info(f"     └─ Estratto{score_info}: {snippet}")

            # Logga tempo di elaborazione
            rag_logger.info(f"\n⏱️Tempo elaborazione: {response.processing_time:.2f}s")

            rag_logger.info(f"Riassunto:\n {response.summary}")

        except KeyboardInterrupt:
            break
        except Exception as e:
            print(f"❌ Errore: {e}")

    print("\n👋 Arrivederci!")


def batch_mode(pipeline, queries_file):
    """Modalità batch da file"""
    try:
        queries_path = Path(queries_file)
        if not queries_path.exists():
            print(f"❌ File non trovato: {queries_file}")
            return

        with open(queries_path, "r", encoding="utf-8") as f:
            queries = [line.strip() for line in f if line.strip()]

        if not queries:
            print("❌ Nessuna query trovata nel file")
            return

        print(f"🔄 Processando {len(queries)} query da {queries_file}")

        responses = pipeline.batch_query(
            queries, top_k=get_default_config().fusion_retrieval.top_k
        )

        # Salva risultati
        output_file = queries_path.with_suffix(".results.txt")
        with open(output_file, "w", encoding="utf-8") as f:
            for i, (query, response) in enumerate(zip(queries, responses), 1):
                f.write(f"Query {i}: {query}\n")
                f.write(f"Risposta: {response.answer}\n")
                f.write(f"Confidence: {response.confidence:.2f}\n")
                f.write(f"Fonti: {', '.join(response.sources)}\n")
                f.write(f"Tempo: {response.processing_time:.2f}s\n")
                f.write("-" * 80 + "\n")

        print(f"✅ Risultati salvati in: {output_file}")

    except Exception as e:
        print(f"❌ Errore nel batch processing: {e}")


def main():
    """Funzione principale"""
    parser = argparse.ArgumentParser(
        description="Sistema RAG Avanzato",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
            Esempi di utilizzo:
            python main.py --docs data/                    # Processa documenti e modalità interattiva
            python main.py --docs data/ --query "Come..."  # Singola query
            python main.py --batch queries.txt             # Batch da file
            python main.py --interactive                   # Solo modalità interattiva (usa cache)
            """,
    )

    parser.add_argument("--docs", nargs="+", help="Percorsi ai documenti da processare")

    parser.add_argument("--query", help="Singola query da eseguire")

    parser.add_argument("--batch", help="File con query per processing batch")

    parser.add_argument(
        "--interactive", action="store_true", help="Modalità interattiva"
    )

    parser.add_argument(
        "--force-reprocess",
        action="store_true",
        help="Forza il riprocessamento anche se esiste cache",
    )

    parser.add_argument(
        "--summary",
        choices=["brief", "comprehensive", "key_points"],
        help="Genera riassunto dei documenti",
    )

    args = parser.parse_args()

    # Setup ambiente
    env_ok, logger = setup_environment()
    if not env_ok:
        sys.exit(1)

    # Inizializza pipeline
    try:
        config = get_default_config()
        pipeline = create_rag_pipeline(config)
        logger.info("Pipeline inizializzata")
    except Exception as e:
        print(f"❌ Errore nell'inizializzazione: {e}")
        sys.exit(1)

    # Processa documenti se specificati
    if args.docs:
        print(f"📁 Processando documenti: {args.docs}")
        if not process_documents(pipeline, args.docs, args.force_reprocess):
            sys.exit(1)

    # Esegui azioni richieste
    if args.query:
        # Singola query
        print(f"🔍 Query: {args.query}")
        try:
            response = pipeline.query(args.query)
            print(f"\n🎯 Risposta:\n{response.answer}")
            if response.sources:
                print(f"\n📚 Fonti: {', '.join(response.sources)}")
        except Exception as e:
            print(f"❌ Errore nella query: {e}")

    elif args.batch:
        # Batch processing
        batch_mode(pipeline, args.batch)

    elif args.summary:
        # Genera riassunto
        try:
            print(f"📄 Generando riassunto ({args.summary})...")
            summary = pipeline.get_document_summary(args.summary)
            print(f"\n📋 Riassunto:\n{summary.answer}")
        except Exception as e:
            print(f"❌ Errore nel riassunto: {e}")

    elif args.interactive or not any([args.query, args.batch, args.summary]):
        # Modalità interattiva (default se nessun'altra azione)
        interactive_mode(pipeline)


if __name__ == "__main__":
    main()