"""Resource mapping cache for metadata stage optimization."""

import polars as pl
from typing import Dict, List, Optional, Any, Tuple
from dataclasses import dataclass
import logging

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class K8sAllocation:
    """
    Immutable K8s allocation data for a resource.

    Represents one K8s workload assignment for a resource.
    Resources can have multiple K8s allocations (2-5 per resource).
    """
    k8s_namespace: str
    k8s_workload: str
    k8s_cluster: str
    k8s_labels: str

    def to_dict(self) -> Dict[str, str]:
        """Convert to dictionary for DataFrame construction."""
        return {
            'k8s_namespace': self.k8s_namespace,
            'k8s_workload': self.k8s_workload,
            'k8s_cluster': self.k8s_cluster,
            'k8s_labels': self.k8s_labels
        }

    @classmethod
    def from_row(cls, row: Dict[str, Any]) -> 'K8sAllocation':
        """Create from DataFrame row."""
        return cls(
            k8s_namespace=row.get('k8s_namespace', ''),
            k8s_workload=row.get('k8s_workload', ''),
            k8s_cluster=row.get('k8s_cluster', ''),
            k8s_labels=row.get('k8s_labels', '')
        )

    def __hash__(self):
        """Enable use in sets for deduplication."""
        return hash((self.k8s_namespace, self.k8s_workload,
                    self.k8s_cluster, self.k8s_labels))


class ResourceMappingCache:
    """
    In-memory cache for resource-to-metadata mappings with K8s allocation support.

    Optimizes multi-month processing by caching metadata allocations from the first
    month and reusing them for subsequent months. Since metadata allocation is
    deterministic (same resource → same metadata across all months), this provides
    4-8x speedup for months 2+.

    Handles multiple K8s allocations per resource (2-5 workloads per resource)
    by accumulating K8s variations while storing base metadata once.

    Cache lifetime: Single run only (cleared on process exit).
    Thread-safety: Not required (single-threaded processing).

    Architecture:
        {
            "provider_name": {
                "AWS:AmazonEC2:instance/instance-123": {
                    "base_metadata": {
                        "resource_id": "instance/app-name-123",
                        "resource/tag:aws:Team": "platform",
                        "region": "us-east-1"
                    },
                    "k8s_allocations": [
                        K8sAllocation(namespace="kube-system", workload="...", ...),
                        K8sAllocation(namespace="default", workload="...", ...)
                    ]
                }
            }
        }
    """

    def __init__(self):
        """Initialize empty cache."""
        self._cache: Dict[str, Dict[str, Dict[str, Any]]] = {}
        self._stats = {
            "cache_hits": 0,
            "cache_misses": 0,
            "resources_cached": 0,
            "partial_hits": 0
        }

    def has_provider_cache(self, provider: str) -> bool:
        """
        Check if provider has cached mappings.

        Args:
            provider: Provider name (e.g., "aws", "azure")

        Returns:
            True if provider cache exists
        """
        return provider in self._cache

    def _separate_k8s_columns(
        self,
        metadata_columns: List[str]
    ) -> Tuple[List[str], List[str]]:
        """
        Separate K8s columns from base metadata columns.

        Args:
            metadata_columns: Full list of metadata column names

        Returns:
            Tuple of (k8s_columns, base_metadata_columns)
        """
        k8s_columns = [
            col for col in metadata_columns
            if col in ['k8s_namespace', 'k8s_workload', 'k8s_cluster', 'k8s_labels']
        ]
        base_columns = [
            col for col in metadata_columns
            if col not in k8s_columns
        ]
        return k8s_columns, base_columns

    def _has_k8s_allocation(self, k8s_alloc: K8sAllocation) -> bool:
        """
        Check if K8s allocation has actual data (not empty/unallocated).

        Args:
            k8s_alloc: K8s allocation to check

        Returns:
            True if has actual allocation, False if empty or unallocated
        """
        return (
            k8s_alloc.k8s_namespace and
            k8s_alloc.k8s_namespace != '__UNALLOCATED__' and
            k8s_alloc.k8s_workload and
            k8s_alloc.k8s_workload != '__UNALLOCATED__'
        )

    def _create_fingerprint_column(self, df: pl.DataFrame) -> pl.DataFrame:
        """
        Create fingerprint column for cache lookup.

        Uses _original_resource_id if present (matches cache population),
        otherwise falls back to resource_id (for backward compatibility).

        Args:
            df: DataFrame with required columns

        Returns:
            DataFrame with _fingerprint column added
        """
        # Determine which resource_id column to use
        resource_id_col = (
            "_original_resource_id"
            if "_original_resource_id" in df.columns
            else "resource_id"
        )

        return df.with_columns([
            pl.concat_str([
                pl.col("cloud_provider"),
                pl.lit(":"),
                pl.col("service"),
                pl.lit(":"),
                pl.col(resource_id_col)
            ]).alias("_fingerprint")
        ])

    def _expand_cache_to_dataframe(self, provider: str) -> pl.DataFrame:
        """
        Convert cache to DataFrame with K8s allocations expanded to rows.

        Each resource with N K8s allocations becomes N rows in the DataFrame.
        Base metadata is duplicated across all rows for the same resource.

        Args:
            provider: Provider name

        Returns:
            DataFrame with cache data (base metadata + K8s allocations expanded)
        """
        cache_rows = []

        for fingerprint, cache_entry in self._cache[provider].items():
            base_metadata = cache_entry['base_metadata']
            k8s_allocations = cache_entry['k8s_allocations']

            # If resource has K8s allocations, create one row per allocation
            if k8s_allocations:
                for k8s_alloc in k8s_allocations:
                    row = {"_fingerprint": fingerprint}
                    row.update(base_metadata)
                    row.update(k8s_alloc.to_dict())
                    cache_rows.append(row)
            else:
                # No K8s allocations - single row with base metadata only
                row = {"_fingerprint": fingerprint}
                row.update(base_metadata)
                cache_rows.append(row)

        return pl.DataFrame(cache_rows) if cache_rows else None

    def lookup_resources(
        self,
        provider: str,
        df: pl.DataFrame,
        hybrid: bool = True
    ) -> tuple[Optional[pl.DataFrame], Optional[pl.DataFrame]]:
        """
        Lookup cached metadata for resources with hybrid mode support.

        Hybrid mode (default): Splits resources into cached and uncached.
        Legacy mode: All-or-nothing (returns None if any resource missing).

        Args:
            provider: Provider name
            df: DataFrame with resources (must have cloud_provider, service, resource_id)
            hybrid: If True, return split of cached/uncached. If False, use all-or-nothing.

        Returns:
            Tuple of (cached_df, uncached_df):
            - cached_df: Resources found in cache with metadata applied (or None)
            - uncached_df: Resources not in cache needing processing (or None)

            In legacy mode (hybrid=False): Returns (None, None) on partial miss
        """
        if df.height == 0:
            return df, None

        if provider not in self._cache:
            self._stats["cache_misses"] += 1
            return None, df

        # Validate required columns
        # Check for _original_resource_id first (matches cache population)
        resource_id_col = (
            "_original_resource_id"
            if "_original_resource_id" in df.columns
            else "resource_id"
        )
        required_cols = ['cloud_provider', 'service', resource_id_col]
        missing_cols = [col for col in required_cols if col not in df.columns]
        if missing_cols:
            logger.warning(f"Cannot use cache - missing columns: {missing_cols}")
            self._stats["cache_misses"] += 1
            return None, df

        # Create fingerprints using helper method
        df_with_fingerprint = self._create_fingerprint_column(df)

        # Convert cache to DataFrame with K8s allocations expanded
        cache_df = self._expand_cache_to_dataframe(provider)

        if cache_df is None or cache_df.height == 0:
            logger.warning(f"Cache exists for {provider} but is empty")
            self._stats["cache_misses"] += 1
            return None, df

        # Left join on fingerprint
        result = df_with_fingerprint.join(
            cache_df,
            on="_fingerprint",
            how="left",
            suffix="_cached"
        )

        # Identify cached columns
        cached_columns = [col for col in cache_df.columns if col != "_fingerprint"]
        if not cached_columns:
            logger.warning(f"No metadata columns in cache for {provider}")
            self._stats["cache_misses"] += 1
            return None, df

        # Determine which column to check for coverage
        first_cached_col = cached_columns[0]
        if first_cached_col + "_cached" in result.columns:
            check_col = first_cached_col + "_cached"
        else:
            check_col = first_cached_col

        missing_count = result[check_col].null_count()

        # Legacy all-or-nothing mode
        if not hybrid and missing_count > 0:
            logger.warning(
                f"Partial cache coverage for {provider}: "
                f"{missing_count}/{df.height} resources not in cache - "
                f"falling back to full processing"
            )
            self._stats["cache_misses"] += 1
            return None, df

        # Hybrid mode: Split cached and uncached
        if missing_count > 0:
            # Split into cached and uncached resources
            cached_mask = pl.col(check_col).is_not_null()

            cached_result = result.filter(cached_mask)
            # Get unique uncached fingerprints (deduplicate since incoming df has 1 row per resource)
            uncached_fingerprints = result.filter(~cached_mask).select("_fingerprint").unique()

            # Get original rows for uncached resources (1 row per fingerprint in input)
            uncached_df = df_with_fingerprint.join(
                uncached_fingerprints,
                on="_fingerprint",
                how="inner"
            ).drop("_fingerprint")

            # Process cached result: coalesce cached columns
            for col in cached_columns:
                if col + "_cached" in cached_result.columns:
                    cached_result = cached_result.with_columns([
                        pl.coalesce([col + "_cached", col]).alias(col)
                    ]).drop(col + "_cached")

            cached_df = cached_result.drop("_fingerprint")

            self._stats["cache_hits"] += 1
            self._stats["partial_hits"] = self._stats.get("partial_hits", 0) + 1

            unique_cached = cached_df["_fingerprint"].n_unique() if "_fingerprint" in cached_df.columns else len(cached_df)
            logger.info(
                f"🔄 Hybrid cache for {provider}: "
                f"{unique_cached:,} cached resources → {cached_df.height:,} rows (with K8s expansions), "
                f"{uncached_df.height:,} need processing"
            )

            return cached_df, uncached_df

        # Full cache hit - all resources found
        for col in cached_columns:
            if col + "_cached" in result.columns:
                result = result.with_columns([
                    pl.coalesce([col + "_cached", col]).alias(col)
                ]).drop(col + "_cached")

        result = result.drop("_fingerprint")

        self._stats["cache_hits"] += 1
        unique_resources = len(self._cache[provider])
        logger.info(
            f"🚀 Full cache HIT for {provider}: {unique_resources:,} resources → "
            f"{result.height:,} rows (with K8s expansions)"
        )

        return result, None

    def _process_cache_chunk(
        self,
        provider: str,
        chunk: pl.DataFrame,
        columns_to_select: List[str],
        k8s_columns: List[str],
        base_columns: List[str],
        fingerprint_resource_id_col: str
    ):
        """
        Process a single chunk for caching with K8s allocation accumulation.

        Separates base metadata (stored once) from K8s allocations (accumulated).
        Maximum 50 lines to comply with CLAUDE.md.

        Args:
            provider: Provider name
            chunk: DataFrame chunk to process
            columns_to_select: All columns needed for caching
            k8s_columns: K8s column names
            base_columns: Non-K8s metadata column names
            fingerprint_resource_id_col: Column for fingerprint generation
        """
        # Use iter_rows only on chunks (acceptable per CLAUDE.md for small chunks)
        for row in chunk.select(columns_to_select).iter_rows(named=True):
            fingerprint = f"{row['cloud_provider']}:{row['service']}:{row[fingerprint_resource_id_col]}"

            # Get or initialize cache entry for this fingerprint
            if fingerprint not in self._cache[provider]:
                self._cache[provider][fingerprint] = {
                    'base_metadata': {},
                    'k8s_allocations': []
                }

            cache_entry = self._cache[provider][fingerprint]

            # Store base metadata (non-K8s columns) - only stored once
            if not cache_entry['base_metadata']:
                cache_entry['base_metadata'] = {
                    col: row[col]
                    for col in base_columns
                    if row[col] is not None
                }

            # Accumulate K8s allocations if present
            if k8s_columns:
                k8s_alloc = K8sAllocation.from_row(row)
                # Only store actual K8s allocations, skip empty/unallocated
                if self._has_k8s_allocation(k8s_alloc):
                    # Use set for automatic deduplication
                    if k8s_alloc not in cache_entry['k8s_allocations']:
                        cache_entry['k8s_allocations'].append(k8s_alloc)

    def cache_provider_mappings(
        self,
        provider: str,
        df: pl.DataFrame,
        metadata_columns: List[str],
        fingerprint_resource_id_col: str = "resource_id"
    ):
        """
        Cache resource mappings with K8s allocation accumulation.

        Separates base metadata from K8s allocations to handle multi-row
        K8s assignments (2-5 allocations per resource).

        Args:
            provider: Provider name
            df: Processed DataFrame with metadata
            metadata_columns: List of metadata column names to cache
            fingerprint_resource_id_col: Column for fingerprint (default: "resource_id")
        """
        if df.height == 0:
            return

        # Validate required columns
        required_cols = ['cloud_provider', 'service', fingerprint_resource_id_col]
        missing_cols = [col for col in required_cols if col not in df.columns]
        if missing_cols:
            logger.warning(f"Cannot cache - missing columns: {missing_cols}")
            return

        if provider not in self._cache:
            self._cache[provider] = {}

        # Validate and separate K8s from base metadata columns
        available_metadata_cols = [col for col in metadata_columns if col in df.columns]
        if not available_metadata_cols:
            logger.warning(f"No metadata columns available to cache for {provider}")
            return

        k8s_columns, base_columns = self._separate_k8s_columns(available_metadata_cols)

        # Build column list for selection
        columns_to_select = required_cols + available_metadata_cols

        # Process in chunks for memory efficiency
        chunk_size = 1_000_000
        for chunk_start in range(0, df.height, chunk_size):
            chunk_end = min(chunk_start + chunk_size, df.height)
            chunk = df.slice(chunk_start, chunk_end - chunk_start)

            self._process_cache_chunk(
                provider, chunk, columns_to_select,
                k8s_columns, base_columns, fingerprint_resource_id_col
            )

            del chunk  # Free memory immediately

        self._stats["resources_cached"] += df.height
        logger.debug(
            f"Cached {df.height:,} resources for {provider} "
            f"({len(base_columns)} base + {len(k8s_columns)} K8s columns)"
        )

    def get_cache_stats(self) -> Dict[str, Any]:
        """
        Get cache statistics for observability.

        Returns:
            Dictionary with cache stats
        """
        total_resources = sum(
            len(resources) for resources in self._cache.values()
        )

        return {
            "providers_cached": len(self._cache),
            "total_resources_in_cache": total_resources,
            "cache_hits": self._stats["cache_hits"],
            "cache_misses": self._stats["cache_misses"],
            "resources_cached": self._stats["resources_cached"],
            "partial_hits": self._stats.get("partial_hits", 0)
        }

    def write_to_disk(self, dataset_name: str, provider: str):
        """
        Write cache to disk for debugging and inspection.

        Creates parquet file in ~/.czbg/cache/metadata/{dataset_name}/{provider}/
        with K8s allocations expanded to rows (same format as cache lookup).

        Args:
            dataset_name: Dataset name for file path
            provider: Provider name (e.g., "aws", "azure")
        """
        if provider not in self._cache or not self._cache[provider]:
            logger.warning(f"No cache data to write for provider {provider}")
            return

        from pathlib import Path

        # Create cache directory
        cache_dir = Path.home() / ".czbg" / "cache" / "metadata" / dataset_name / provider.lower()
        cache_dir.mkdir(parents=True, exist_ok=True)

        # Use the same expansion helper as cache lookup for consistency
        cache_df = self._expand_cache_to_dataframe(provider)

        if cache_df is None or cache_df.height == 0:
            logger.warning(f"No cache data to expand for {provider}")
            return

        # Write to parquet
        cache_file = cache_dir / "cache.parquet"
        cache_df.write_parquet(cache_file)

        # Calculate stats for logging
        unique_resources = len(self._cache[provider])
        total_k8s_allocations = sum(
            len(entry['k8s_allocations'])
            for entry in self._cache[provider].values()
            if entry['k8s_allocations']
        )

        logger.info(
            f"Wrote cache to disk: {cache_file} "
            f"({unique_resources:,} resources → {cache_df.height:,} rows, "
            f"{total_k8s_allocations:,} K8s allocations)"
        )

    def clear(self):
        """Clear all cached mappings."""
        self._cache.clear()
        self._stats = {
            "cache_hits": 0,
            "cache_misses": 0,
            "resources_cached": 0,
            "partial_hits": 0
        }
        logger.debug("Resource mapping cache cleared")
