#!/usr/bin/env python3
"""
Enable Email/Password authentication for a Firebase project via Identity Platform API.

This script handles the FULL initialization flow:
1. Links a billing account (required for Identity Platform)
2. Initializes Identity Platform if not already initialized
3. Enables Email/Password authentication

Usage:
    python enable_email_auth.py <project_id>
    
Example:
    python enable_email_auth.py symph-ai-teacher-assistant

Requires:
    - google-auth, requests (pip install google-auth requests)
    - Service account with Firebase Admin + Billing Admin roles
"""

import json
import os
import sys
import time
from google.oauth2 import service_account
from google.auth.transport.requests import Request
import requests

# Default billing account to use (Abaca - known to be open)
DEFAULT_BILLING_ACCOUNT = "billingAccounts/00446F-8B8C0A-26B94E"

def get_credentials(credentials_path=None):
    """Get OAuth credentials."""
    if credentials_path is None:
        credentials_path = os.environ.get(
            "GOOGLE_APPLICATION_CREDENTIALS",
            os.path.expanduser("~/.config/gcloud/aria-firebase-admin.json")
        )
    
    if not os.path.exists(credentials_path):
        print(f"❌ Credentials file not found: {credentials_path}")
        return None
    
    SCOPES = [
        'https://www.googleapis.com/auth/cloud-platform',
        'https://www.googleapis.com/auth/firebase',
        'https://www.googleapis.com/auth/cloud-billing'
    ]
    
    try:
        credentials = service_account.Credentials.from_service_account_file(
            credentials_path, scopes=SCOPES
        )
        credentials.refresh(Request())
        return credentials
    except Exception as e:
        print(f"❌ Failed to get credentials: {e}")
        return None


def find_open_billing_account(headers):
    """Find an open billing account to use."""
    url = "https://cloudbilling.googleapis.com/v1/billingAccounts"
    resp = requests.get(url, headers=headers)
    
    if not resp.ok:
        return DEFAULT_BILLING_ACCOUNT
    
    accounts = resp.json().get("billingAccounts", [])
    open_accounts = [a for a in accounts if a.get("open", False)]
    
    if open_accounts:
        # Prefer "Symph" related accounts
        for acc in open_accounts:
            name = acc.get("displayName", "").lower()
            if "symph" in name:
                return acc["name"]
        return open_accounts[0]["name"]
    
    return DEFAULT_BILLING_ACCOUNT


def ensure_billing_enabled(project_id, headers, billing_account=None):
    """Ensure billing is enabled for the project."""
    # Check current billing status
    url = f"https://cloudbilling.googleapis.com/v1/projects/{project_id}/billingInfo"
    resp = requests.get(url, headers=headers)
    
    if resp.ok:
        data = resp.json()
        if data.get("billingEnabled"):
            print(f"✅ Billing already enabled")
            return True
    
    # Find an open billing account if not specified
    if billing_account is None:
        billing_account = find_open_billing_account(headers)
    
    print(f"📋 Linking billing account: {billing_account}")
    
    # Link billing account
    resp = requests.put(url, headers=headers, json={"billingAccountName": billing_account})
    
    if resp.ok:
        data = resp.json()
        if data.get("billingEnabled"):
            print(f"✅ Billing enabled")
            return True
        else:
            print(f"⚠️ Billing account linked but not enabled (account may be closed)")
            return False
    else:
        print(f"❌ Failed to link billing: {resp.status_code}")
        print(f"   {resp.text[:200]}")
        return False


def initialize_identity_platform(project_id, headers):
    """Initialize Identity Platform for the project."""
    url = f"https://identitytoolkit.googleapis.com/v2/projects/{project_id}/identityPlatform:initializeAuth"
    resp = requests.post(url, headers=headers, json={})
    
    if resp.ok:
        print(f"✅ Identity Platform initialized")
        return True
    elif resp.status_code == 409:
        # Already initialized
        print(f"✅ Identity Platform already initialized")
        return True
    elif "BILLING_NOT_ENABLED" in resp.text:
        print(f"❌ Billing not enabled - cannot initialize Identity Platform")
        return False
    else:
        print(f"❌ Failed to initialize Identity Platform: {resp.status_code}")
        print(f"   {resp.text[:200]}")
        return False


def enable_email_auth(project_id, credentials_path=None):
    """Enable Email/Password authentication for a Firebase project."""
    
    credentials = get_credentials(credentials_path)
    if not credentials:
        return False
    
    print(f"🔑 Got credentials for project: {project_id}")
    
    headers = {
        "Authorization": f"Bearer {credentials.token}",
        "Content-Type": "application/json"
    }
    
    # Step 1: Ensure billing is enabled
    print(f"\n📋 Step 1: Checking billing...")
    if not ensure_billing_enabled(project_id, headers):
        print("\n💡 Tip: Link an OPEN billing account to the project in GCP Console")
        return False
    
    # Step 2: Initialize Identity Platform
    print(f"\n🚀 Step 2: Initializing Identity Platform...")
    if not initialize_identity_platform(project_id, headers):
        return False
    
    # Brief pause for propagation
    time.sleep(2)
    
    # Step 3: Enable email/password auth
    print(f"\n📧 Step 3: Enabling Email/Password auth...")
    url = f"https://identitytoolkit.googleapis.com/admin/v2/projects/{project_id}/config?updateMask=signIn.email"
    payload = {
        "signIn": {
            "email": {
                "enabled": True,
                "passwordRequired": True
            }
        }
    }
    
    resp = requests.patch(url, headers=headers, json=payload)
    
    if resp.ok:
        print(f"\n🎉 Email/Password authentication ENABLED for {project_id}!")
        return True
    else:
        print(f"❌ Failed to enable email auth: {resp.status_code}")
        print(f"   {resp.text[:200]}")
        return False


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python enable_email_auth.py <project_id>")
        print("Example: python enable_email_auth.py symph-ai-teacher-assistant")
        sys.exit(1)
    
    project_id = sys.argv[1]
    success = enable_email_auth(project_id)
    sys.exit(0 if success else 1)
