"""
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 normalize_anthropic_service_to_model(service: str) -> str:
    """
    Convert Anthropic Service column values to model tag values.

    Args:
        service: Service column value (e.g., "Claude Sonnet 4", "Claude Opus 4.1")

    Returns:
        Normalized model name (e.g., "claude-sonnet-4", "claude-opus-4-1")
    """
    if not service:
        return "unknown"

    # Convert to lowercase and replace spaces with hyphens
    normalized = service.lower().replace(" ", "-")

    # Handle version numbers with dots (4.1 -> 4-1)
    normalized = normalized.replace(".", "-")

    return normalized


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,
    job_uuid: 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)
        job_uuid: Unique job identifier to ensure uniqueness across jobs

    Returns:
        12-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}"
    if job_uuid:
        seed_data += f":{job_uuid}"

    uid = hashlib.md5(seed_data.encode()).hexdigest()[:12]
    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 generate_deterministic_account_id(provider: str, uid: str) -> str:
    """Generate deterministic 6-character account ID from provider and UID.

    Args:
        provider: Cloud provider name (AWS, Azure, GCP)
        uid: 12-character deterministic UID

    Returns:
        Account ID in format: {normalized-provider}-{6-char-uid}
    """
    normalized_provider = normalize(provider or "")
    account_suffix = uid[:6].upper()
    return f"{normalized_provider}-{account_suffix}"


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,
        "cloud_account": 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 (subscriptionName and resourceGroupName use azure: prefix)
        "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,
        # Anthropic-specific tags
        "resource/tag:anthropic_cz:model": 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",
    job_uuid: str = None,
) -> 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
        job_uuid: Unique job identifier to ensure uniqueness across jobs
        
    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"),
            job_uuid,
        )
        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,
        "resource/tag:anthropic_cz:model": [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(resource_type or "")),
                pl.lit("-"),
                pl.col("uid").str.slice(0, 12)
            ]).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("cloud_account"),
            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(resource_type or "")),
                pl.lit("-"),
                pl.col("uid").str.slice(0, 12)
            ]).alias("name")
        ]).with_columns([
            pl.concat_str([
                pl.lit(f"{resource_type}/"),
                pl.col("name")
            ]).alias("resource_id"),
            pl.col("project").alias("cloud_account"),
            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"])

    elif provider == "Anthropic":
        # Map Service column to model tag
        model_tag_value = normalize_anthropic_service_to_model(service)

        df = df.with_columns([
            pl.col("uid").map_elements(
                lambda uid: generate_deterministic_account_id(provider, uid),
                return_dtype=pl.Utf8
            ).alias("cloud_account"),
            pl.concat_str([
                pl.lit(f"{resource_type}/"),
                pl.lit(normalize(resource_type or "")),
                pl.lit("-"),
                pl.col("uid").str.slice(0, 12)
            ]).alias("resource_id"),
            # Set the model tag based on the Service column
            pl.lit(model_tag_value).alias("resource/tag:anthropic_cz:model")
        ])

    else:  # AWS and other providers
        df = df.with_columns([
            pl.col("uid").map_elements(
                lambda uid: generate_deterministic_account_id(provider, uid),
                return_dtype=pl.Utf8
            ).alias("cloud_account"),
            pl.concat_str([
                # Use : separator for AWS, / for others (temporary workaround)
                pl.when(pl.col("cloud_provider") == "AWS")
                  .then(pl.lit(f"{resource_type}:"))
                  .otherwise(pl.lit(f"{resource_type}/")),
                pl.lit(normalize(resource_type or "")),
                pl.lit("-"),
                pl.col("uid").str.slice(0, 12)
            ]).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",
    job_uuid: str = None,
) -> 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
        job_uuid: Unique job identifier to ensure uniqueness across jobs

    Returns:
        DataFrame with all generated resources
    """
    if filtered_df.is_empty():
        schema = _get_consistent_schema()
        return pl.DataFrame([], schema=schema)

    # Sort filtered DataFrame for deterministic ordering across runs
    # This ensures resource_index is assigned consistently to the same reference rows
    # Fill nulls with empty string for sorting (nulls can't be sorted)
    filtered_df = filtered_df.with_columns([
        pl.col('usage_family').fill_null(""),
        pl.col('operation').fill_null(""),
        pl.col('usage_type').fill_null("")
    ]).sort([
        'cloud_provider', 'service', 'resource_type',
        'usage_family', 'operation', 'usage_type'
    ])

    # Expand each reference row into resource_count copies
    # Using repeat_by + explode for vectorized expansion
    expanded_df = filtered_df.with_columns([
        pl.lit(resource_count).alias("_repeat_count")
    ]).with_columns([
        pl.col(col).repeat_by("_repeat_count")
        for col in filtered_df.columns
    ]).drop("_repeat_count").explode(filtered_df.columns)

    # Add resource_index for each reference row group
    # Must maintain sorted order for deterministic index assignment
    # Re-sort after explode to ensure consistent ordering for window function
    group_cols = ['cloud_provider', 'service', 'resource_type', 'usage_family', 'operation', 'usage_type']
    expanded_df = expanded_df.sort(group_cols).with_columns([
        pl.int_range(0, pl.len()).over(group_cols).alias("resource_index")
    ])

    # Generate UIDs vectorized using struct + map_elements
    expanded_df = expanded_df.with_columns([
        pl.struct([
            'cloud_provider', 'service', 'resource_type', 'resource_index',
            'usage_family', 'operation', 'usage_type'
        ]).map_elements(
            lambda row: generate_deterministic_uid(
                org,
                row['cloud_provider'],
                row['service'],
                row['resource_type'],
                row['resource_index'],
                row['usage_family'],
                row['operation'],
                row['usage_type'],
                job_uuid
            ),
            return_dtype=pl.Utf8
        ).alias("uid")
    ])

    # Initialize all provider tag columns as null
    expanded_df = expanded_df.with_columns([
        pl.lit(None).alias("resource/tag:azure:subscriptionName"),
        pl.lit(None).alias("resource/tag:azure:resourceGroupName"),
        pl.lit(None).alias("resource/tag:gcp_cz:project_name"),
        pl.lit(None).alias("resource/tag:gcp_cz:resource_name"),
        pl.lit(None).alias("resource/tag:gcp_cz:ancestry_display_name"),
        pl.lit(None).alias("resource/tag:gcp:compute.googleapis.com/machine_spec"),
        pl.lit(None).alias("resource/tag:anthropic_cz:model"),
        pl.lit("Usage").alias("line_item_type")
    ])

    # Process provider-specific logic in batches
    # Azure resources
    azure_mask = pl.col("cloud_provider") == "Azure"
    expanded_df = expanded_df.with_columns([
        pl.when(azure_mask).then(
            pl.col("uid").map_elements(get_deterministic_subscription, return_dtype=pl.Utf8)
        ).otherwise(None).alias("_azure_subscription"),
        pl.when(azure_mask).then(
            pl.concat_str([pl.lit("rg-"), pl.col("uid").str.slice(0, 4)])
        ).otherwise(None).alias("_azure_rg"),
        pl.when(azure_mask).then(
            pl.concat_str([
                pl.col("resource_type").str.to_lowercase().str.replace_all(" ", "-", literal=True).str.replace_all("/", "-", literal=True),
                pl.lit("-"),
                pl.col("uid").str.slice(0, 12)
            ])
        ).otherwise(None).alias("_azure_name")
    ]).with_columns([
        pl.when(azure_mask).then(
            pl.concat_str([
                pl.lit("/subscriptions/"),
                pl.col("_azure_subscription"),
                pl.lit("/resourceGroups/"),
                pl.col("_azure_rg"),
                pl.lit("/providers/"),
                pl.col("service"),
                pl.lit("/"),
                pl.col("resource_type"),
                pl.lit("/"),
                pl.col("_azure_name")
            ])
        ).alias("_azure_resource_id"),
        pl.col("_azure_subscription").alias("resource/tag:azure:subscriptionName"),
        pl.col("_azure_rg").alias("resource/tag:azure:resourceGroupName")
    ])

    # GCP resources
    gcp_mask = pl.col("cloud_provider") == "GCP"
    expanded_df = expanded_df.with_columns([
        pl.when(gcp_mask).then(
            pl.col("uid").map_elements(get_deterministic_project, return_dtype=pl.Utf8)
        ).otherwise(None).alias("_gcp_project"),
        pl.when(gcp_mask).then(
            pl.col("uid").map_elements(get_deterministic_ancestry, return_dtype=pl.Utf8)
        ).otherwise(None).alias("_gcp_ancestry"),
        pl.when(gcp_mask).then(
            pl.concat_str([
                pl.col("resource_type").str.to_lowercase().str.replace_all(" ", "-", literal=True).str.replace_all("/", "-", literal=True),
                pl.lit("-"),
                pl.col("uid").str.slice(0, 12)
            ])
        ).otherwise(None).alias("_gcp_name")
    ]).with_columns([
        pl.when(gcp_mask).then(
            pl.concat_str([pl.col("resource_type"), pl.lit("/"), pl.col("_gcp_name")])
        ).alias("_gcp_resource_id"),
        pl.col("_gcp_project").alias("resource/tag:gcp_cz:project_name"),
        pl.when(gcp_mask).then(
            pl.concat_str([pl.col("resource_type"), pl.lit("/"), pl.col("_gcp_name")])
        ).otherwise(pl.col("resource/tag:gcp_cz:resource_name")).alias("resource/tag:gcp_cz:resource_name"),
        pl.col("_gcp_ancestry").alias("resource/tag:gcp_cz:ancestry_display_name")
    ])

    # Anthropic resources
    anthropic_mask = pl.col("cloud_provider") == "Anthropic"
    expanded_df = expanded_df.with_columns([
        pl.when(anthropic_mask).then(
            pl.col("service").map_elements(normalize_anthropic_service_to_model, return_dtype=pl.Utf8)
        ).otherwise(pl.col("resource/tag:anthropic_cz:model")).alias("resource/tag:anthropic_cz:model"),
        pl.when(anthropic_mask).then(
            pl.concat_str([
                pl.col("resource_type"),
                pl.lit("/"),
                pl.col("resource_type").str.to_lowercase().str.replace_all(" ", "-", literal=True).str.replace_all("/", "-", literal=True),
                pl.lit("-"),
                pl.col("uid").str.slice(0, 12)
            ])
        ).alias("_anthropic_resource_id")
    ])

    # AWS resources with : separator (temporary workaround)
    aws_mask = pl.col("cloud_provider") == "AWS"
    expanded_df = expanded_df.with_columns([
        pl.when(aws_mask).then(
            pl.concat_str([
                pl.col("resource_type"),
                pl.lit(":"),  # Use : for AWS
                pl.col("resource_type").str.to_lowercase().str.replace_all(" ", "-", literal=True).str.replace_all("/", "-", literal=True),
                pl.lit("-"),
                pl.col("uid").str.slice(0, 12)
            ])
        ).alias("_aws_resource_id")
    ])

    # Other providers with / separator
    other_mask = ~(azure_mask | gcp_mask | anthropic_mask | aws_mask)
    expanded_df = expanded_df.with_columns([
        pl.when(other_mask).then(
            pl.concat_str([
                pl.col("resource_type"),
                pl.lit("/"),  # Use / for others
                pl.col("resource_type").str.to_lowercase().str.replace_all(" ", "-", literal=True).str.replace_all("/", "-", literal=True),
                pl.lit("-"),
                pl.col("uid").str.slice(0, 12)
            ])
        ).alias("_other_resource_id")
    ])

    # Combine resource_id from all providers
    expanded_df = expanded_df.with_columns([
        pl.coalesce([
            pl.col("_azure_resource_id"),
            pl.col("_gcp_resource_id"),
            pl.col("_anthropic_resource_id"),
            pl.col("_aws_resource_id"),
            pl.col("_other_resource_id")
        ]).alias("resource_id")
    ])

    # Cloud account handling
    expanded_df = expanded_df.with_columns([
        pl.when(azure_mask).then(pl.col("_azure_subscription"))
        .when(gcp_mask).then(pl.col("_gcp_project"))
        .otherwise(
            pl.struct(['cloud_provider', 'uid']).map_elements(
                lambda row: generate_deterministic_account_id(row['cloud_provider'], row['uid']),
                return_dtype=pl.Utf8
            )
        ).alias("cloud_account")
    ])

    # Drop temporary columns
    temp_cols = [c for c in expanded_df.columns if c.startswith("_")]
    expanded_df = expanded_df.drop(temp_cols + ["resource_index", "uid"])

    # Apply consistent schema
    schema = _get_consistent_schema()
    result_df = expanded_df.cast(schema)

    return result_df