#!/usr/bin/env python3
"""Experiment: which MongoDB connection method works in strict/cloud environments?

Tests PyMongo vs mongosh (shell) vs other approaches. Run locally and on EC2
to find what works where PyMongo SSL fails.

Usage: uv run python scripts/mongodb_connection_experiment.py
"""

import json
import os
import subprocess
import sys
from pathlib import Path

# ensure we load .env from project root
root = Path(__file__).parent.parent
env_file = root / ".env"
if env_file.exists():
    from dotenv import load_dotenv
    load_dotenv(env_file)

URI = os.getenv("PT_RESEARCH_DATABASE_URI")
if not URI:
    print("SKIP: PT_RESEARCH_DATABASE_URI not set")
    sys.exit(0)


def test_pymongo() -> tuple[bool, str]:
    """Test PyMongo with certifi + optional tlsAllowInvalidCertificates."""
    try:
        import certifi
        from pymongo import MongoClient
        opts = {
            "serverSelectionTimeoutMS": 5000,
            "connectTimeoutMS": 5000,
            "tlsCAFile": certifi.where(),
        }
        if os.getenv("PT_RESEARCH_TLS_SKIP_VERIFY", "").strip() in ("1", "true", "yes"):
            opts["tlsAllowInvalidCertificates"] = True
        client = MongoClient(URI, **opts)
        db = client.get_default_database()
        count = db["Operator"].count_documents({})
        return True, f"PyMongo OK: {count} operators"
    except Exception as e:
        return False, f"PyMongo FAIL: {e}"


def test_mongosh() -> tuple[bool, str]:
    """Test mongosh via subprocess - uses mongosh's own SSL stack."""
    try:
        # mongosh can add tlsAllowInvalidCertificates to URI
        uri = URI
        if os.getenv("PT_RESEARCH_TLS_SKIP_VERIFY", "").strip() in ("1", "true", "yes"):
            sep = "&" if "?" in uri else "?"
            uri = f"{uri}{sep}tlsAllowInvalidCertificates=true"
        result = subprocess.run(
            ["mongosh", uri, "--eval", "db.Operator.countDocuments({})", "--quiet"],
            capture_output=True,
            text=True,
            timeout=15,
        )
        if result.returncode == 0:
            count = result.stdout.strip()
            return True, f"mongosh OK: {count} operators"
        return False, f"mongosh FAIL: {result.stderr or result.stdout or result.returncode}"
    except FileNotFoundError:
        return False, "mongosh FAIL: mongosh not installed"
    except subprocess.TimeoutExpired:
        return False, "mongosh FAIL: timeout"
    except Exception as e:
        return False, f"mongosh FAIL: {e}"


def test_mongosh_json() -> tuple[bool, str]:
    """Test mongosh with JSON output for structured data."""
    try:
        uri = URI
        if os.getenv("PT_RESEARCH_TLS_SKIP_VERIFY", "").strip() in ("1", "true", "yes"):
            sep = "&" if "?" in uri else "?"
            uri = f"{uri}{sep}tlsAllowInvalidCertificates=true"
        result = subprocess.run(
            [
                "mongosh", uri,
                "--eval", "JSON.stringify(db.Operator.find().limit(2).toArray())",
                "--quiet",
            ],
            capture_output=True,
            text=True,
            timeout=15,
        )
        if result.returncode == 0:
            data = json.loads(result.stdout.strip())
            names = [d.get("name", "?") for d in data]
            return True, f"mongosh JSON OK: sample operators {names}"
        return False, f"mongosh JSON FAIL: {result.stderr or result.stdout}"
    except FileNotFoundError:
        return False, "mongosh JSON FAIL: mongosh not installed"
    except json.JSONDecodeError as e:
        return False, f"mongosh JSON FAIL: parse error {e}"
    except Exception as e:
        return False, f"mongosh JSON FAIL: {e}"


def test_mongo_legacy() -> tuple[bool, str]:
    """Test legacy mongo shell (mongo) if mongosh not available."""
    try:
        uri = URI
        if os.getenv("PT_RESEARCH_TLS_SKIP_VERIFY", "").strip() in ("1", "true", "yes"):
            sep = "&" if "?" in uri else "?"
            uri = f"{uri}{sep}tlsAllowInvalidCertificates=true"
        result = subprocess.run(
            ["mongo", uri, "--eval", "db.Operator.count()", "--quiet"],
            capture_output=True,
            text=True,
            timeout=15,
        )
        if result.returncode == 0:
            return True, f"mongo (legacy) OK: {result.stdout.strip()} operators"
        return False, f"mongo FAIL: {result.stderr or result.stdout}"
    except FileNotFoundError:
        return False, "mongo FAIL: mongo not installed"
    except Exception as e:
        return False, f"mongo FAIL: {e}"


def main():
    print("=" * 60)
    print("MongoDB connection experiment")
    print("=" * 60)
    print(f"URI: {URI[:50]}... (hidden)")
    print(f"PT_RESEARCH_TLS_SKIP_VERIFY: {os.getenv('PT_RESEARCH_TLS_SKIP_VERIFY', '(not set)')}")
    print()

    results = []
    for name, fn in [
        ("PyMongo (certifi + optional skip)", test_pymongo),
        ("mongosh (subprocess)", test_mongosh),
        ("mongosh JSON output", test_mongosh_json),
        ("mongo legacy (if installed)", test_mongo_legacy),
    ]:
        ok, msg = fn()
        results.append((name, ok, msg))
        icon = "✓" if ok else "✗"
        print(f"  {icon} {name}: {msg}")

    print()
    working = [r for r in results if r[1]]
    if working:
        print("Working methods:", [r[0] for r in working])
    else:
        print("No method worked. Check URI, network, and that mongosh is installed if trying shell.")
    print()
    return 0 if any(r[1] for r in results) else 1


if __name__ == "__main__":
    sys.exit(main())
