"""Main metadata processor orchestrator."""

import polars as pl
from typing import Dict, Any, List
import logging

from .config import MetadataConfig
from .services.random_service import DeterministicRandomService
from .services.tags_service import TagsAllocationService
from .services.four_group_service import FourGroupAllocationService
from .services.region_service import RegionAllocationService
from .services.account_service import AccountAllocationService
from .services.correlated_service import CorrelatedAllocationService
from .services.azure_resource_id_service import AzureResourceIdService
from .services.gcp_tags_service import GcpTagsService
from .services.key_variant_resolver import KeyVariantResolver
from .services.value_variant_selector import ValueVariantSelector
from .services.variant_column_manager import VariantColumnManager
from .services.variant_aware_expression_service import VariantAwareExpressionService
from .services.k8s_labels_aggregation_service import K8sLabelsAggregationService
from .services.resource_id_only_service import ResourceIdOnlyService
from .services.expression_builder import ExpressionBuilder
from .services.expression_cache import ExpressionCache
from .services.resource_mapping_cache import ResourceMappingCache

logger = logging.getLogger(__name__)


class MetadataProcessor:
    """Main orchestrator for metadata allocation processing."""

    def __init__(self, config: MetadataConfig):
        """
        Initialize processor with configuration.

        Args:
            config: Complete metadata configuration with providers and keys
        """
        self.config = config

        # Initialize core services
        self.random_service = DeterministicRandomService(config.dataset_name)
        self.azure_resource_id_service = AzureResourceIdService()
        self.gcp_tags_service = GcpTagsService()

        # Initialize variant services
        self.key_variant_resolver = KeyVariantResolver(config.key_variants)
        self.value_variant_selector = ValueVariantSelector(self.random_service)
        self.variant_column_manager = VariantColumnManager()

        # Initialize allocation services with dependency injection
        self.tags_service = TagsAllocationService(self.random_service)
        self.variant_expression_service = VariantAwareExpressionService(
            self.tags_service,
            self.key_variant_resolver
        )
        self.four_group_service = FourGroupAllocationService(
            self.random_service,
            self.key_variant_resolver
        )
        self.region_service = RegionAllocationService(self.random_service)
        self.account_service = AccountAllocationService(
            self.random_service,
            self.azure_resource_id_service
        )
        self.correlated_service = CorrelatedAllocationService(
            self.random_service,
            self.azure_resource_id_service,
            self.key_variant_resolver
        )
        self.k8s_aggregation_service = K8sLabelsAggregationService(
            self.random_service,
            self.key_variant_resolver
        )
        self.resource_id_only_service = ResourceIdOnlyService(
            self.random_service
        )

        # NEW: Expression building and caching
        self.expression_builder = ExpressionBuilder(
            random_service=self.random_service,
            tags_service=self.tags_service,
            four_group_service=self.four_group_service,
            correlated_service=self.correlated_service,
            k8s_aggregation_service=self.k8s_aggregation_service,
            region_service=self.region_service,
            account_service=self.account_service,
            variant_expression_service=self.variant_expression_service,
            resource_id_only_service=self.resource_id_only_service
        )

        self.expression_cache = ExpressionCache(self.expression_builder)
        self.resource_cache = ResourceMappingCache()

        logger.info(f"Initialized MetadataProcessor for dataset: {config.dataset_name}")

    def process_provider_month(
        self,
        provider: str,
        month: str,
        df: pl.DataFrame
    ) -> pl.DataFrame:
        """
        Process metadata allocation for a single provider-month.

        Uses hybrid ResourceMappingCache for optimal performance:
        - Full cache hit: Skip processing entirely
        - Partial cache hit: Process only new resources, combine with cached
        - Cache miss: Process all resources (first month)

        Uses ExpressionCache for efficient expression reuse across months.

        Args:
            provider: Provider name (e.g., "aws", "azure", "gcp")
            month: Month identifier (e.g., "2025-01")
            df: DataFrame with resources from this provider-month

        Returns:
            DataFrame with metadata allocation applied
        """
        if df.height == 0:
            logger.warning(f"Empty DataFrame for {provider}/{month}")
            return df

        if not self.config.has_provider(provider):
            raise ValueError(
                f"Provider '{provider}' not found in configuration"
            )

        provider_config = self.config.get_provider_config(provider)

        # Try hybrid cache lookup (months 2+)
        if self.resource_cache.has_provider_cache(provider):
            cached_df, uncached_df = self.resource_cache.lookup_resources(
                provider, df, hybrid=True
            )

            # Full cache hit - all resources found
            if uncached_df is None or uncached_df.height == 0:
                return cached_df

            # Partial cache hit - process uncached resources
            logger.info(
                f"Processing {uncached_df.height:,} new resources for {provider}/{month}"
            )
            processed_uncached = self._process_dataframe(
                provider, provider_config, uncached_df
            )

            # Add newly processed resources to cache
            # Use _original_resource_id directly to avoid creating a copy
            metadata_columns = self._identify_metadata_columns(processed_uncached)
            self.resource_cache.cache_provider_mappings(
                provider, processed_uncached, metadata_columns,
                fingerprint_resource_id_col="_original_resource_id"
            )

            # Write updated cache to disk for debugging
            self.resource_cache.write_to_disk(self.config.dataset_name, provider)

            # Drop temporary column before combining
            processed_uncached = processed_uncached.drop(
                ["_original_resource_id", "_preprocessed_resource_id"]
            )

            # Combine cached and newly processed
            result = pl.concat([cached_df, processed_uncached])

            logger.info(
                f"✅ Completed {provider}/{month}: "
                f"{cached_df.height:,} cached + {processed_uncached.height:,} processed"
            )

            return result

        # Cache miss - process all resources (month 1)
        logger.info(
            f"Processing {provider}/{month}: {df.height:,} resources"
        )

        # Memory optimization: Process in chunks if dataset is large
        chunk_size = 5_000_000  # 5M rows per chunk
        if df.height > chunk_size:
            logger.info(f"Large dataset detected ({df.height:,} rows), processing in chunks of {chunk_size:,}")
            result = self._process_dataframe_chunked(provider, provider_config, df, chunk_size)
        else:
            result = self._process_dataframe(provider, provider_config, df)

        # Cache the mappings for future months
        # Use _original_resource_id directly to avoid creating a copy of the DataFrame
        metadata_columns = self._identify_metadata_columns(result)
        self.resource_cache.cache_provider_mappings(
            provider, result, metadata_columns, fingerprint_resource_id_col="_original_resource_id"
        )

        # Write cache to disk for debugging (first month only)
        self.resource_cache.write_to_disk(self.config.dataset_name, provider)

        # Drop temporary columns after caching
        result = result.drop(["_original_resource_id", "_preprocessed_resource_id"])

        logger.info(
            f"✅ Completed {provider}/{month}: {result.height:,} resources "
            f"(cached {len(metadata_columns)} metadata columns)"
        )

        # Memory debug: log memory before return
        try:
            import psutil
            import os
            process = psutil.Process(os.getpid())
            mem_gb = process.memory_info().rss / 1024**3
            logger.debug(f"Memory before return: {mem_gb:.2f} GB (process RSS)")
        except:
            pass

        return result

    def _process_dataframe(
        self,
        provider: str,
        provider_config: 'ProviderConfig',
        df: pl.DataFrame
    ) -> pl.DataFrame:
        """
        Apply full metadata processing to a DataFrame.

        Args:
            provider: Provider name
            provider_config: Provider configuration
            df: DataFrame to process

        Returns:
            Processed DataFrame with metadata applied (includes temporary columns)
        """
        # Phase 1: Preprocessing - Drop resource type prefixes and preserve state
        df = df.with_columns([
            pl.col("resource_id").alias("_original_resource_id")
        ])
        df = self._drop_resource_type_prefixes(df)
        df = df.with_columns([
            pl.col("resource_id").alias("_preprocessed_resource_id")
        ])

        # Phase 2: Get and apply expressions
        logger.debug(f"Getting expressions for {provider}")
        expression_groups = self.expression_cache.get_expressions(
            provider, provider_config, df
        )
        resource_id_exprs = expression_groups['resource_id']
        batch_exprs = expression_groups['batch']
        logger.debug(
            f"Got {len(resource_id_exprs)} resource_id + {len(batch_exprs)} batch "
            f"expressions for {provider}"
        )

        try:
            # Step 1: Apply resource_id expressions sequentially (chain injection)
            result = df
            for i, expr in enumerate(resource_id_exprs):
                logger.debug(f"Applying resource_id expression {i+1}/{len(resource_id_exprs)}")
                result = result.with_columns(expr)

            # Step 2: Apply batch expressions together (memory-efficient)
            if batch_exprs:
                logger.debug(f"Applying {len(batch_exprs)} batch expressions to {result.height} rows")
                result = result.with_columns(batch_exprs)

            logger.debug(
                f"Successfully applied expressions, result has {result.height} rows, "
                f"{len(result.columns)} columns"
            )
        except Exception as e:
            logger.error(f"Failed applying expressions for {provider}: {e}")
            logger.error(f"DataFrame columns: {df.columns}")
            logger.error(
                f"Number of expressions: {len(resource_id_exprs)} resource_id + "
                f"{len(batch_exprs)} batch"
            )
            raise RuntimeError(
                f"Failed to apply metadata expressions for {provider}: {e}"
            ) from e

        # Phase 3: Restoration - Restore original if no injection occurred
        result = result.with_columns([
            pl.when(pl.col("resource_id") == pl.col("_preprocessed_resource_id"))
            .then(pl.col("_original_resource_id"))  # No injection → restore
            .otherwise(pl.col("resource_id"))        # Injection occurred → keep
            .alias("resource_id")
        ])

        # Azure post-processing
        if provider.lower() == "azure" and "resource_id" in result.columns:
            result = self._apply_azure_resource_id_updates(
                result, provider_config.correlated_groups, provider
            )

        # GCP post-processing - sync resource_name with resource_id changes
        if provider.lower() in ['gcp', 'google cloud platform', 'google cloud', 'google']:
            if "resource_id" in result.columns and "resource/tag:gcp_cz:resource_name" in result.columns:
                result = self.gcp_tags_service.sync_resource_name_with_resource_id(result)

        return result

    def _process_dataframe_chunked(
        self,
        provider: str,
        provider_config: 'ProviderConfig',
        df: pl.DataFrame,
        chunk_size: int
    ) -> pl.DataFrame:
        """
        Process DataFrame in chunks to reduce memory usage.

        Writes each processed chunk to temporary parquet files to avoid
        accumulating all chunks in memory. Final result is read back
        using Polars' efficient scan/concat.

        Args:
            provider: Provider name
            provider_config: Provider configuration
            df: DataFrame to process
            chunk_size: Number of rows per chunk

        Returns:
            Processed DataFrame with all chunks combined
        """
        import gc
        import tempfile
        from pathlib import Path

        total_rows = df.height
        num_chunks = (total_rows + chunk_size - 1) // chunk_size
        logger.info(f"Processing {total_rows:,} rows in {num_chunks} chunks")

        # Create temporary directory for chunk files
        # Use disk-based location (not /tmp which is often tmpfs/RAM-backed)
        # Prefer ~/.czbg/temp for actual disk storage
        import os
        base_temp_dir = Path.home() / ".czbg" / "temp"
        base_temp_dir.mkdir(parents=True, exist_ok=True)
        temp_dir = tempfile.mkdtemp(prefix="metadata_chunks_", dir=str(base_temp_dir))
        chunk_files = []

        try:
            for i in range(num_chunks):
                start_idx = i * chunk_size
                end_idx = min((i + 1) * chunk_size, total_rows)

                logger.debug(f"Processing chunk {i+1}/{num_chunks}: rows {start_idx:,} to {end_idx:,}")

                # Extract chunk
                chunk = df.slice(start_idx, end_idx - start_idx)

                # Process chunk
                processed_chunk = self._process_dataframe(provider, provider_config, chunk)

                # Write chunk to temporary parquet file (memory-efficient)
                chunk_file = Path(temp_dir) / f"chunk_{i:04d}.parquet"
                processed_chunk.write_parquet(chunk_file)
                chunk_files.append(chunk_file)

                # Free memory immediately
                del chunk, processed_chunk
                gc.collect()

                logger.debug(f"Completed chunk {i+1}/{num_chunks}, written to {chunk_file.name}")

            # Free original DataFrame before combining - no longer needed
            # Must do this before concat to reduce peak memory
            logger.debug(f"Freeing original DataFrame ({df.height:,} rows) before combining chunks")
            del df
            gc.collect()

            # Combine all chunks using incremental concatenation
            logger.info(f"Combining {len(chunk_files)} processed chunks from disk")

            # Read and combine chunks incrementally to minimize peak memory
            # Instead of holding all chunks then concat, we concat as we go
            result = None
            for i, chunk_file in enumerate(chunk_files):
                logger.debug(f"Reading and combining chunk {i+1}/{len(chunk_files)}")
                chunk_df = pl.read_parquet(chunk_file)

                if result is None:
                    result = chunk_df
                else:
                    # Concat with existing result
                    result = pl.concat([result, chunk_df])
                    # Free the chunk immediately
                    del chunk_df
                    gc.collect()

            logger.info(f"Chunked processing complete: {result.height:,} rows")
            return result

        finally:
            # Clean up temporary files
            import shutil
            try:
                shutil.rmtree(temp_dir)
                logger.debug(f"Cleaned up temporary chunk files in {temp_dir}")
            except Exception as e:
                logger.warning(f"Failed to clean up temporary chunk files: {e}")

    def _identify_metadata_columns(self, df: pl.DataFrame) -> List[str]:
        """
        Identify metadata columns that should be cached.

        Metadata columns include:
        - resource/tag:* columns
        - k8s_labels
        - k8s_namespace
        - k8s_workload (for cross-month consistency)
        - k8s_cluster (for cross-month consistency)
        - resource_id (may be modified by injection)
        - region (special case column)
        - cloud_account (special case column)

        Args:
            df: Processed DataFrame

        Returns:
            List of metadata column names
        """
        metadata_columns = []

        for col in df.columns:
            # Tag columns
            if col.startswith("resource/tag:"):
                metadata_columns.append(col)
            # K8s columns - cache ALL K8s metadata for consistency
            elif col in ["k8s_labels", "k8s_namespace", "k8s_workload", "k8s_cluster"]:
                metadata_columns.append(col)
            # Resource ID (may be modified)
            elif col == "resource_id":
                metadata_columns.append(col)
            # Special case columns
            elif col in ["region", "cloud_account"]:
                metadata_columns.append(col)

        return metadata_columns

    def _apply_azure_resource_id_updates(
        self,
        df: pl.DataFrame,
        correlated_groups: list,
        provider: str
    ) -> pl.DataFrame:
        """
        Apply Azure resource ID updates for correlated groups containing resourceGroupName.

        Args:
            df: DataFrame with applied expressions
            correlated_groups: List of correlated group configurations
            provider: Provider name

        Returns:
            DataFrame with updated resource IDs
        """
        result_df = df

        for group_config in correlated_groups:
            # Check if this correlated group contains resourceGroupName
            for key_name in group_config.keys:
                if key_name.lower() == "resourcegroupname":
                    # Use the metadata tag column as the source for resource group names
                    tag_column = f"resource/tag:{provider}:{key_name}"
                    if tag_column in result_df.columns:
                        logger.debug(f"Applying Azure resource ID update for {tag_column}")
                        result_df = self.azure_resource_id_service.update_resource_group_in_resource_id(
                            result_df, tag_column
                        )
                    break

        return result_df

    def get_processing_summary(self, provider: str) -> Dict[str, Any]:
        """
        Get summary of what will be processed for a provider.

        Args:
            provider: Provider name

        Returns:
            Dictionary with processing summary
        """
        if not self.config.has_provider(provider):
            return {
                "provider": provider,
                "configured": False,
                "error": f"Provider '{provider}' not found in configuration"
            }

        provider_config = self.config.get_provider_config(provider)

        # Categorize keys by processing type
        region_keys = []
        account_keys = []
        tags_only_keys = []
        k8s_labels_only_keys = []
        resource_id_only_keys = []
        four_group_keys = []

        for key_config in provider_config.keys:
            if self.region_service.is_region_key(key_config.key_name):
                region_keys.append({
                    "key": key_config.key_name,
                    "values_count": len(key_config.values)
                })
            elif self.account_service.is_account_key(key_config.key_name):
                account_keys.append({
                    "key": key_config.key_name,
                    "values_count": len(key_config.values)
                })
            elif key_config.tags_only:
                tags_only_keys.append({
                    "key": key_config.key_name,
                    "values_count": len(key_config.values)
                })
            elif key_config.k8s_labels_only:
                k8s_labels_only_keys.append({
                    "key": key_config.key_name,
                    "values_count": len(key_config.values)
                })
            elif key_config.resource_id_only:
                resource_id_only_keys.append({
                    "key": key_config.key_name,
                    "values_count": len(key_config.values)
                })
            else:
                four_group_keys.append({
                    "key": key_config.key_name,
                    "values_count": len(key_config.values)
                })

        return {
            "provider": provider,
            "configured": True,
            "total_keys": len(provider_config.keys),
            "processing_types": {
                "region_keys": region_keys,
                "account_keys": account_keys,
                "tags_only_keys": tags_only_keys,
                "k8s_labels_only_keys": k8s_labels_only_keys,
                "resource_id_only_keys": resource_id_only_keys,
                "four_group_keys": four_group_keys
            }
        }

    def validate_dataframe_requirements(self, df: pl.DataFrame) -> list[str]:
        """
        Validate that DataFrame has required columns for processing.

        Args:
            df: DataFrame to validate

        Returns:
            List of validation errors (empty if valid)
        """
        errors = []

        if df.height == 0:
            return errors  # Empty DataFrame is valid

        # Check for basic required columns
        required_columns = ["cloud_provider", "service", "resource_id"]
        missing_columns = [col for col in required_columns if col not in df.columns]

        if missing_columns:
            errors.append(f"DataFrame missing required columns: {missing_columns}")

        return errors

    def _drop_resource_type_prefixes(self, df: pl.DataFrame) -> pl.DataFrame:
        """
        Drop resource type prefix from resource_id.

        Handles provider-specific formats:
        - Azure: Special path-based format with subscriptions
        - All others (AWS, GCP, MongoDB, etc.): type/id format where only id part is modified

        For non-Azure providers:
        - Format: type/id (e.g., "service-usage/service-usage-9a7d4245ae0e")
        - Preserves type part before "/"
        - Drops prefix only from id part after "/"
        - Result: type/cleaned-id (e.g., "service-usage/usage-9a7d4245ae0e")

        Args:
            df: DataFrame with resource_id column

        Returns:
            DataFrame with resource type prefixes removed from id part only
        """
        return df.with_columns([
            pl.when(pl.col("cloud_provider") == "Azure")
            .then(
                # Azure: /subscriptions/.../vm-1a2b → /subscriptions/.../1a2b
                self._drop_azure_resource_type()
            )
            .otherwise(
                # All other providers: type/id format
                # Only modify id part (after /), preserve type part (before /)
                self._drop_standard_resource_type()
            )
            .alias("resource_id")
        ])

    def _drop_azure_resource_type(self) -> pl.Expr:
        """
        Drop Azure resource type (last segment after final slash).

        Example: /subscriptions/.../vm-1a2b → /subscriptions/.../1a2b
        """
        # Split on '/', get last segment, split on '-', take rest
        segments = pl.col("resource_id").str.split("/")
        last_segment = segments.list.get(-1)

        # Check if last segment has hyphen
        return (
            pl.when(last_segment.str.contains("-"))
            .then(
                # Drop first part before hyphen in last segment
                pl.concat_str([
                    segments.list.slice(0, segments.list.len() - 1).list.join("/"),
                    pl.lit("/"),
                    last_segment.str.split("-").list.slice(1).list.join("-")
                ])
            )
            .otherwise(pl.col("resource_id"))
        )

    def _drop_standard_resource_type(self) -> pl.Expr:
        """
        Drop resource type prefix from standard type/id format.

        Applies to all non-Azure providers (AWS, GCP, MongoDB, Anthropic, etc.)

        Format:
        - AWS: type:id (e.g., "queue:queue-abc123")
        - Others: type/id (e.g., "service-usage/service-usage-9a7d4245ae0e")

        Where:
        - type: Resource type (preserved unchanged)
        - id: Resource identifier with optional type prefix to remove

        Process:
        1. Split on ':' (AWS) or '/' (others) to separate type from id
        2. Drop prefix from id part (first segment before hyphen in id)
        3. Reconstruct as type:cleaned-id (AWS) or type/cleaned-id (others)

        Examples:
        - AWS: "queue:queue-abc123" → "queue:abc123"
        - GCP: "service-usage/service-usage-9a7d4245ae0e" → "service-usage/usage-9a7d4245ae0e"
        - GCP: "instances/inst-1a2b" → "instances/1a2b"
        """
        # Split on both : and / to handle AWS (: separator) and other providers (/ separator)
        segments_colon = pl.col("resource_id").str.split(":")
        segments_slash = pl.col("resource_id").str.split("/")

        # Use : split for AWS, / split for others
        type_part = pl.when(pl.col("cloud_provider") == "AWS") \
            .then(segments_colon.list.get(0)) \
            .otherwise(segments_slash.list.get(0))

        id_part = pl.when(pl.col("cloud_provider") == "AWS") \
            .then(segments_colon.list.get(-1)) \
            .otherwise(segments_slash.list.get(-1))

        # Reconstruct with appropriate separator
        separator = pl.when(pl.col("cloud_provider") == "AWS") \
            .then(pl.lit(":")) \
            .otherwise(pl.lit("/"))

        # Drop prefix from id part (first segment before hyphen)
        # Check if id part has hyphen to split
        return (
            pl.when(id_part.str.contains("-"))
            .then(
                # Drop first part before hyphen in id segment
                pl.concat_str([
                    type_part,
                    separator,
                    id_part.str.split("-").list.slice(1).list.join("-")
                ])
            )
            .otherwise(pl.col("resource_id"))
        )

