"""
Resource generation for billing data.

Extracted from the existing pipeline with deterministic UID generation
and consistent schema handling.
"""

import hashlib
import random
from typing import Dict
import polars as pl


GCP_ANCESTRY_NAMES = [
    "cloudloom.com/prod/data-bq",
    "cloudloom.com/dev/data-bq",
    "cloudloom.com/stg/data-bq",
    "cloudloom.com/prod/infra-core",
    "cloudloom.com/dev/infra-core",
    "cloudloom.com/stg/infra-core",
    "cloudloom.com/prod/shared-services",
    "cloudloom.com/dev/shared-services",
    "cloudloom.com/stg/shared-services",
]

AZURE_SUBSCRIPTIONS = [
    "Enterprise IT Subscription",
    "Data Analytics Platform",
    "Development Environment",
    "Production Workloads",
    "Shared Services Infrastructure",
]

GCP_PROJECT_NAMES = [
    "Global Operations",
    "Customer Data Platform",
    "Research Computing",
    "Marketing Analytics",
    "Core Infrastructure",
]


def normalize(s: str) -> str:
    """Normalize string for resource naming."""
    return s.lower().replace(" ", "-").replace("/", "-")


def generate_deterministic_uid(
    org: str,
    provider: str,
    service: str,
    resource_type: str,
    resource_index: int,
    usage_family: str = None,
    operation: str = None,
    usage_type: str = None,
) -> str:
    """
    Generate deterministic UID that's consistent across months.

    Args:
        org: Organization name
        provider: Cloud provider (AWS, Azure, GCP)
        service: Service name (e.g., AmazonEC2, microsoft.compute)
        resource_type: Resource type (e.g., instances, storageAccounts)
        resource_index: Index of resource within service/type combination
        usage_family: Usage family (e.g., Compute Instance, Database Instance)
        operation: Operation name (e.g., RunInstances, GetObject)
        usage_type: Specific usage type (e.g., BoxUsage:c4.8xlarge)

    Returns:
        8-character deterministic UID
    """
    # Include full metadata combination for granular resource IDs
    seed_data = f"{org}:{provider}:{service}:{resource_type}:{resource_index}"

    if usage_family:
        seed_data += f":{usage_family}"
    if operation:
        seed_data += f":{operation}"
    if usage_type:
        seed_data += f":{usage_type}"

    uid = hashlib.md5(seed_data.encode()).hexdigest()[:8]
    return uid


def get_deterministic_subscription(uid: str) -> str:
    """Get deterministic Azure subscription based on UID."""
    subscription_index = int(uid[:2], 16) % len(AZURE_SUBSCRIPTIONS)
    return AZURE_SUBSCRIPTIONS[subscription_index]


def get_deterministic_project(uid: str) -> str:
    """Get deterministic GCP project based on UID."""
    project_index = int(uid[:2], 16) % len(GCP_PROJECT_NAMES)
    return GCP_PROJECT_NAMES[project_index]


def get_deterministic_ancestry(uid: str) -> str:
    """Get deterministic GCP ancestry based on UID."""
    ancestry_index = int(uid[2:4], 16) % len(GCP_ANCESTRY_NAMES)
    return GCP_ANCESTRY_NAMES[ancestry_index]


def _get_consistent_schema() -> Dict[str, pl.DataType]:
    """
    Generate consistent schema for all generated resources.

    Ensures all provider tag columns are typed as String to prevent
    "type String is incompatible with expected type Null" errors during
    concatenation operations.

    Returns:
        Dictionary mapping column names to Polars data types
    """
    # Base schema for core resource fields
    schema = {
        "resource_id": pl.Utf8,
        "cloud_provider": pl.Utf8,
        "service": pl.Utf8,
        "resource_type": pl.Utf8,
        "usage_family": pl.Utf8,
        "operation": pl.Utf8,
        "usage_type": pl.Utf8,
        "line_item_type": pl.Utf8,
    }

    # Add all provider tag columns as String type
    # This stays in sync with the tag initialization in generate_resources()
    provider_tag_columns = {
        # Azure-specific tags
        "resource/tag:azure:subscriptionName": pl.Utf8,
        "resource/tag:azure:resourceGroupName": pl.Utf8,
        # GCP-specific tags
        "resource/tag:gcp_cz:project_name": pl.Utf8,
        "resource/tag:gcp_cz:resource_name": pl.Utf8,
        "resource/tag:gcp_cz:ancestry_display_name": pl.Utf8,
        "resource/tag:gcp:compute.googleapis.com/machine_spec": pl.Utf8,
    }

    # Merge provider tags into schema
    schema.update(provider_tag_columns)

    return schema


def get_gcp_machine_spec(
    ref_df: pl.DataFrame, cloud_provider: str, service: str, resource_type: str
) -> str:
    """Get GCP machine specification from reference data."""
    matches = (
        ref_df.filter(
            (pl.col("cloud_provider") == cloud_provider)
            & (pl.col("service") == service)
            & (pl.col("resource_type") == resource_type)
            & (pl.col("instance_type").is_not_null())
        )
        .select("instance_type")
        .to_series()
        .unique()
        .to_list()
    )
    return random.choice(matches) if matches else ""


def generate_resources(
    ref_row: Dict[str, str],
    resource_count: int,
    ref_df: pl.DataFrame,
    org: str = "default",
) -> pl.DataFrame:
    """
    Generate resources based on reference data row using vectorized operations.
    
    Args:
        ref_row: Dictionary containing reference data fields
        resource_count: Number of resources to generate
        ref_df: Full reference DataFrame for context
        org: Organization name for deterministic UIDs
        
    Returns:
        DataFrame with generated resources
    """
    provider = ref_row.get("cloud_provider")
    service = ref_row.get("service")
    resource_type = ref_row.get("resource_type")

    # Vectorized UID generation
    uids = []
    for i in range(resource_count):
        uid = generate_deterministic_uid(
            org,
            provider,
            service,
            resource_type,
            i,
            ref_row.get("usage_family"),
            ref_row.get("operation"),
            ref_row.get("usage_type"),
        )
        uids.append(uid)

    # Create base DataFrame with repeated values
    df_data = {
        "resource_index": list(range(resource_count)),
        "uid": uids,
        "cloud_provider": [provider] * resource_count,
        "service": [service] * resource_count,
        "resource_type": [resource_type] * resource_count,
        "usage_family": [ref_row.get("usage_family")] * resource_count,
        "operation": [ref_row.get("operation")] * resource_count,
        "usage_type": [ref_row.get("usage_type")] * resource_count,
        "line_item_type": ["Usage"] * resource_count,
    }
    
    # Initialize all provider tag columns as null
    df_data.update({
        "resource/tag:azure:subscriptionName": [None] * resource_count,
        "resource/tag:azure:resourceGroupName": [None] * resource_count,
        "resource/tag:gcp_cz:project_name": [None] * resource_count,
        "resource/tag:gcp_cz:resource_name": [None] * resource_count,
        "resource/tag:gcp_cz:ancestry_display_name": [None] * resource_count,
        "resource/tag:gcp:compute.googleapis.com/machine_spec": [None] * resource_count,
    })
    
    # Create DataFrame for vectorized operations
    df = pl.DataFrame(df_data)
    
    # Vectorized provider-specific processing
    if provider == "Azure":
        df = df.with_columns([
            # Generate Azure resource IDs and tags using vectorized operations
            pl.col("uid").map_elements(
                lambda uid: get_deterministic_subscription(uid), 
                return_dtype=pl.Utf8
            ).alias("subscription"),
            pl.col("uid").map_elements(
                lambda uid: f"rg-{uid[:4]}", 
                return_dtype=pl.Utf8
            ).alias("rg_name"),
            pl.concat_str([
                pl.lit(normalize(service or "")),
                pl.lit("-"),
                pl.lit(normalize(resource_type or "")),
                pl.lit("-"),
                pl.col("uid").str.slice(0, 4)
            ]).alias("name")
        ]).with_columns([
            pl.concat_str([
                pl.lit("/subscriptions/"),
                pl.col("subscription"),
                pl.lit("/resourceGroups/"),
                pl.col("rg_name"),
                pl.lit(f"/providers/{service}/{resource_type}/"),
                pl.col("name")
            ]).alias("resource_id"),
            pl.col("subscription").alias("resource/tag:azure:subscriptionName"),
            pl.col("rg_name").alias("resource/tag:azure:resourceGroupName")
        ]).drop(["subscription", "rg_name", "name"])
        
    elif provider == "GCP":
        # Get machine spec for GCP Compute Engine
        machine_spec = None
        if service == "Compute Engine":
            machine_spec = get_gcp_machine_spec(ref_df, provider, service, resource_type or "")
        
        df = df.with_columns([
            # Generate GCP resource IDs and tags using vectorized operations
            pl.col("uid").map_elements(
                lambda uid: get_deterministic_project(uid), 
                return_dtype=pl.Utf8
            ).alias("project"),
            pl.col("uid").map_elements(
                lambda uid: get_deterministic_ancestry(uid), 
                return_dtype=pl.Utf8
            ).alias("ancestry"),
            pl.concat_str([
                pl.lit(normalize(service or "")),
                pl.lit("-"),
                pl.lit(normalize(resource_type or "")),
                pl.lit("-"),
                pl.col("uid").str.slice(0, 4)
            ]).alias("name")
        ]).with_columns([
            pl.concat_str([
                pl.lit(f"{resource_type}/"),
                pl.col("name")
            ]).alias("resource_id"),
            pl.col("project").alias("resource/tag:gcp_cz:project_name"),
            pl.concat_str([
                pl.lit(f"{resource_type}/"),
                pl.col("name")
            ]).alias("resource/tag:gcp_cz:resource_name"),
            pl.col("ancestry").alias("resource/tag:gcp_cz:ancestry_display_name")
        ])
        
        # Add machine spec if available
        if machine_spec:
            df = df.with_columns(
                pl.lit(machine_spec).alias("resource/tag:gcp:compute.googleapis.com/machine_spec")
            )
            
        df = df.drop(["project", "ancestry", "name"])
        
    else:  # AWS and other providers
        df = df.with_columns([
            pl.concat_str([
                pl.lit(f"{resource_type}/"),
                pl.lit(normalize(service or "")),
                pl.lit("-"),
                pl.lit(normalize(resource_type or "")),
                pl.lit("-"),
                pl.col("uid").str.slice(0, 4)
            ]).alias("resource_id")
        ])

    # Drop temporary columns and ensure schema consistency
    df = df.drop(["resource_index", "uid"])
    
    # Apply consistent schema and collect if it's a lazy frame
    schema = _get_consistent_schema()
    result_df = df.cast(schema)
    
    # Only collect if it's a lazy frame
    if hasattr(result_df, 'collect'):
        return result_df.collect()
    else:
        return result_df


def generate_resources_batch(
    filtered_df: pl.DataFrame,
    resource_count: int,
    ref_df: pl.DataFrame,
    org: str = "default",
) -> pl.DataFrame:
    """
    Generate resources for multiple reference rows using vectorized operations.
    
    Args:
        filtered_df: DataFrame with multiple reference data rows
        resource_count: Number of resources to generate per row
        ref_df: Full reference DataFrame for context
        org: Organization name for deterministic UIDs
        
    Returns:
        DataFrame with all generated resources
    """
    if filtered_df.is_empty():
        schema = _get_consistent_schema()
        return pl.DataFrame([], schema=schema)
    
    all_resources = []
    
    # Process each reference row but collect results for efficient concatenation
    for row in filtered_df.iter_rows(named=True):
        resources = generate_resources(row, resource_count, ref_df, org)
        all_resources.append(resources)
    
    # Single concatenation operation instead of incremental appends
    if all_resources:
        return pl.concat(all_resources, how="vertical")
    else:
        schema = _get_consistent_schema()
        return pl.DataFrame([], schema=schema)