import argparse
import json
import os
from infer_db.api import infer_relationships
from infer_db.utils.report_formatter import save_schema_to_json

def main():
    parser = argparse.ArgumentParser(description="Infer Database Relationships (FKs) automatically.")
    
    # Connection Args
    parser.add_argument("--host", required=True, help="Database Host")
    parser.add_argument("--dbname", required=True, help="Database Name")
    parser.add_argument("--user", required=True, help="Database User")
    parser.add_argument("--password", help="Database Password (optional if env var set)")
    parser.add_argument("--port", default="5432", help="Database Port")
    
    # Output Args
    parser.add_argument("--output", "-o", default="schema_map.json", help="Output JSON file path")
    
    args = parser.parse_args()

    # Build Config
    db_config = {
        "host": args.host,
        "dbname": args.dbname,
        "user": args.user,
        "password": args.password or os.getenv("DB_PASSWORD"),
        "port": args.port
    }

    if not db_config["password"]:
        print("Error: Password must be provided via --password or DB_PASSWORD env var.")
        return

    # Run Inference (Using your Library!)
    print(f"🚀 Connecting to {args.dbname} at {args.host}...")
    try:
        report = infer_relationships(db_config)
        
        # Save Output
        save_schema_to_json(report, args.output)
        print(f"✅ Done! Report saved to: {args.output}")
        
    except Exception as e:
        print(f"❌ Critical Error: {e}")

if __name__ == "__main__":
    main()