#!/usr/bin/env python3
"""
Create audience_members table in BigQuery for Preset Audience migration.

This script creates the audience_members table for one or more properties.
Location: customer-360-profile.client_{property_id}.audience_members

Usage:
    python create_audience_members_table.py [property_id]
    
    If property_id is not provided, creates tables for all properties.
"""

import os
import sys
import logging
from datetime import datetime

# Add system-api to path (now we're inside system-api, so just add current directory)
# system-api/migration/scripts/create_audience_members_table.py -> ../../../
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../../..'))

from connectors.bigquery.bq import BigQuery
from connectors.firebase.firebase import Firebase

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# Table schema for audience_members
AUDIENCE_MEMBERS_SCHEMA = [
    {
        "name": "audience_id",
        "type": "STRING",
        "mode": "REQUIRED",
        "description": "Audience ID from Firebase"
    },
    {
        "name": "user_pseudo_id",
        "type": "STRING",
        "mode": "REQUIRED",
        "description": "User pseudo ID"
    },
    {
        "name": "added_at",
        "type": "DATETIME",
        "mode": "REQUIRED",
        "description": "Timestamp when user was added to audience"
    },
    {
        "name": "removed_at",
        "type": "DATETIME",
        "mode": "NULLABLE",
        "description": "Timestamp when user was removed from audience"
    },
    {
        "name": "is_active",
        "type": "BOOL",
        "mode": "REQUIRED",
        "description": "Whether the user is currently active in the audience"
    }
]

PARTITION_FIELD = "added_at"
CLUSTER_FIELDS = ["audience_id", "user_pseudo_id"]


def create_table_for_property(property_id: str, bq: BigQuery) -> bool:
    """
    Create audience_members table for a specific property.
    
    Args:
        property_id: Property ID (e.g., "12345")
        bq: BigQuery connector instance
        
    Returns:
        True if successful, False otherwise
    """
    try:
        dataset_id = f"client_{property_id}"
        table_id = "audience_members"
        
        logger.info(f"Creating table {dataset_id}.{table_id} for property {property_id}")
        
        bq.create_table_from_json_schema(
            dataset_id=dataset_id,
            table_id=table_id,
            schema_json=AUDIENCE_MEMBERS_SCHEMA,
            partition_field=PARTITION_FIELD,
            cluster_fields=CLUSTER_FIELDS
        )
        
        logger.info(f"✅ Successfully created table {dataset_id}.{table_id}")
        return True
        
    except Exception as e:
        logger.error(f"❌ Failed to create table for property {property_id}: {e}")
        return False


def get_all_properties(fb: Firebase) -> list:
    """
    Get list of all property IDs from Firebase.
    
    Args:
        fb: Firebase connector instance
        
    Returns:
        List of property IDs
    """
    try:
        accounts = fb.db.reference().child("account").get(shallow=True)
        if accounts:
            return list(accounts.keys())
        return []
    except Exception as e:
        logger.error(f"Failed to get properties from Firebase: {e}")
        return []


def main():
    """Main function to create audience_members tables."""
    # Initialize connectors
    try:
        bq = BigQuery()
        fb = Firebase(host=os.environ.get("FIREBASE_HOST"))
    except Exception as e:
        logger.error(f"Failed to initialize connectors: {e}")
        sys.exit(1)
    
    # Get property ID from command line or use all properties
    if len(sys.argv) > 1:
        property_ids = [sys.argv[1]]
        logger.info(f"Creating table for property: {property_ids[0]}")
    else:
        property_ids = get_all_properties(fb)
        logger.info(f"Creating tables for {len(property_ids)} properties")
    
    if not property_ids:
        logger.warning("No properties found")
        sys.exit(1)
    
    # Create tables
    success_count = 0
    failed_count = 0
    
    for property_id in property_ids:
        if create_table_for_property(property_id, bq):
            success_count += 1
        else:
            failed_count += 1
    
    # Summary
    logger.info(f"\n{'='*50}")
    logger.info(f"Summary:")
    logger.info(f"  ✅ Success: {success_count}")
    logger.info(f"  ❌ Failed: {failed_count}")
    logger.info(f"{'='*50}")
    
    if failed_count > 0:
        sys.exit(1)


if __name__ == "__main__":
    main()