#!/usr/bin/env python3
"""
Enable Firebase Storage for a Firebase project programmatically.

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

Prerequisites:
    - google-auth, requests (pip install google-auth requests)
    - Service account with:
      - Firebase Admin role
      - Billing Account User role (to link billing if needed)
    - Cloud Billing API enabled on service account's project

What this does:
    1. Links billing account to project (required for GCS)
    2. Creates App Engine app (creates default bucket)
    3. Adds bucket to Firebase Storage
    4. After this, you can deploy storage.rules with Firebase CLI
"""

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

def enable_storage(project_id: str, credentials_path: str = None) -> bool:
    """Enable Firebase Storage for a project."""
    
    if credentials_path is None:
        credentials_path = 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 False
    
    # Get OAuth token
    SCOPES = [
        'https://www.googleapis.com/auth/cloud-platform',
        'https://www.googleapis.com/auth/cloud-billing',
        'https://www.googleapis.com/auth/firebase'
    ]
    try:
        credentials = service_account.Credentials.from_service_account_file(
            credentials_path, scopes=SCOPES
        )
        credentials.refresh(Request())
        token = credentials.token
    except Exception as e:
        print(f"❌ Failed to get OAuth token: {e}")
        return False
    
    headers = {
        "Authorization": f"Bearer {token}",
        "Content-Type": "application/json"
    }
    
    bucket_name = f"{project_id}.appspot.com"
    
    # Step 1: Check if billing is linked
    print(f"🔍 Checking billing for {project_id}...")
    billing_url = f"https://cloudbilling.googleapis.com/v1/projects/{project_id}/billingInfo"
    resp = requests.get(billing_url, headers=headers)
    
    if resp.status_code == 200:
        billing_info = resp.json()
        if not billing_info.get('billingEnabled'):
            print("📋 Linking billing account...")
            # Get first available billing account
            accounts_resp = requests.get(
                "https://cloudbilling.googleapis.com/v1/billingAccounts",
                headers=headers
            )
            if accounts_resp.status_code == 200:
                accounts = accounts_resp.json().get('billingAccounts', [])
                if accounts:
                    billing_account = accounts[0]['name']
                    link_resp = requests.put(
                        billing_url,
                        headers=headers,
                        json={"billingAccountName": billing_account}
                    )
                    if link_resp.status_code == 200:
                        print(f"✅ Billing linked: {billing_account}")
                    else:
                        print(f"⚠️ Could not link billing: {link_resp.text[:200]}")
    
    # Step 2: Check if storage bucket exists
    print(f"🔍 Checking storage bucket: {bucket_name}...")
    bucket_url = f"https://storage.googleapis.com/storage/v1/b/{bucket_name}"
    resp = requests.get(bucket_url, headers=headers)
    
    if resp.status_code == 404:
        print("📦 Creating App Engine (initializes storage bucket)...")
        # Create App Engine app
        ae_url = "https://appengine.googleapis.com/v1/apps"
        ae_resp = requests.post(ae_url, headers=headers, json={
            "id": project_id,
            "locationId": "us-central"
        })
        
        if ae_resp.status_code in [200, 409]:
            print("⏳ Waiting for bucket to be created...")
            # Wait for bucket to be created
            for i in range(12):  # Wait up to 60 seconds
                time.sleep(5)
                resp = requests.get(bucket_url, headers=headers)
                if resp.status_code == 200:
                    print(f"✅ Storage bucket created: {bucket_name}")
                    break
            else:
                print(f"⚠️ Bucket creation timed out, but may still complete")
        else:
            print(f"⚠️ App Engine creation: {ae_resp.status_code}")
            print(ae_resp.text[:200])
    else:
        print(f"✅ Storage bucket already exists: {bucket_name}")
    
    # Step 3: Add bucket to Firebase Storage
    print("🔥 Adding bucket to Firebase Storage...")
    fb_url = f"https://firebasestorage.googleapis.com/v1beta/projects/{project_id}/buckets/{bucket_name}:addFirebase"
    resp = requests.post(fb_url, headers=headers, json={})
    
    if resp.status_code == 200:
        print(f"✅ Firebase Storage enabled for {project_id}")
        print(f"\n📋 Next step: Deploy storage rules with:")
        print(f"   firebase deploy --only storage --project {project_id}")
        return True
    elif resp.status_code == 409:
        print(f"✅ Firebase Storage already enabled for {project_id}")
        return True
    else:
        print(f"❌ Failed to enable Firebase Storage: {resp.status_code}")
        print(resp.text[:300])
        return False

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