main.py

"""
CLI interface for the NL2SQL Pipeline.

Usage:
    # Interactive mode
    python main.py --db clinic.db

    # Single query
    python main.py --db clinic.db --query "How many patients do we have?"

    # Run evaluation suite
    python evaluate.py
"""

import argparse
import json
import os
import sys
from datetime import datetime

from pipeline import NL2SQLPipeline, PipelineConfig


def interactive_loop(pipe: NL2SQLPipeline, log_path: str | None):
    """Run an interactive NL2SQL session."""
    results_log = []
    print("\n" + "=" * 60)
    print("NL2SQL Pipeline — Interactive Mode")
    print("Ask business questions in plain English.")
    print("Type 'quit' to exit.")
    print("=" * 60 + "\n")

    while True:
        try:
            question = input("Q: ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nExiting.")
            break

        if not question or question.lower() in ("quit", "exit", "q"):
            break

        result = pipe.query(question)
        print(f"\n{result}\n")

        results_log.append({
            "question": result.question,
            "sql": result.cleaned_sql,
            "success": result.success,
            "row_count": result.execution_result.get("row_count", 0),
            "retries": result.retries_used,
            "timestamp": datetime.now().isoformat(),
        })

    # Save log
    if results_log and log_path:
        os.makedirs(os.path.dirname(log_path) if os.path.dirname(log_path) else ".", exist_ok=True)
        with open(log_path, "w") as f:
            json.dump(results_log, f, indent=2)
        print(f"Query log saved to {log_path}")


def main():
    parser = argparse.ArgumentParser(description="NL2SQL Pipeline")
    parser.add_argument("--db", type=str, default="clinic.db", help="Path to SQLite database")
    parser.add_argument("--ddl", type=str, help="Path to DDL file (alternative to --db)")
    parser.add_argument("--query", type=str, help="Single query (non-interactive)")
    parser.add_argument("--log", type=str, default="logs/queries.json", help="Query log path")
    parser.add_argument("--model", type=str, default=None, help="Ollama model name")
    args = parser.parse_args()

    config = PipelineConfig()
    if args.model:
        config.model_name = args.model
    config.db_path = args.db

    pipe = NL2SQLPipeline(config)

    if args.ddl:
        pipe.initialize(ddl_path=args.ddl)
    elif os.path.exists(args.db):
        pipe.initialize(db_path=args.db)
    else:
        print(f"ERROR: Database not found at {args.db}")
        print("Run: python create_sample_db.py")
        sys.exit(1)

    if args.query:
        result = pipe.query(args.query)
        print(f"\n{result}")
    else:
        interactive_loop(pipe, args.log)


if __name__ == "__main__":
    main()