"""
Local Cache Manager for reference data operations.

Handles all local cache operations including add, flush, and change tracking.
Follows Single Responsibility Principle from CLAUDE.md.
"""

import hashlib
import json
from datetime import datetime
from pathlib import Path
from typing import Optional

import polars as pl

from src.log_context import get_logger
from .models import CacheInfo, FilterCriteria, FlushResult, FlushPreview


class LocalCacheManager:
    """
    Manages local reference data cache operations.

    Single responsibility: Local file operations and change tracking.
    """

    def __init__(self, cache_path: Optional[Path] = None):
        """Initialize local cache manager with cache file path."""
        self.logger = get_logger(__name__)

        # Default cache location
        if cache_path is None:
            cache_path = Path.home() / ".czbg" / "cache" / "reference_data" / "reference.parquet"

        self.cache_path = cache_path
        self.cache_path.parent.mkdir(parents=True, exist_ok=True)

        # Local manifest for change tracking
        self.local_manifest_path = self.cache_path.parent / "local_manifest.json"

    def add_records(self, new_data: pl.DataFrame) -> int:
        """
        Add records to local cache, merging with existing data.

        Handles schema harmonization to prevent width mismatch errors.

        Args:
            new_data: DataFrame with new records to add

        Returns:
            Number of new combinations added
        """
        try:
            # Load existing data
            if self.cache_path.exists():
                existing_df = pl.read_parquet(self.cache_path)
            else:
                # Create empty DataFrame with modern schema including instance_type
                existing_df = pl.DataFrame(schema=self._get_standard_schema())

            # Harmonize schemas before processing
            existing_df, new_data = self._harmonize_schemas(existing_df, new_data)

            # Updated key fields to include instance_type
            key_fields = ["cloud_provider", "service", "usage_family",
                         "resource_type", "operation", "usage_type", "instance_type"]

            # Get available key fields (handle legacy data gracefully)
            available_key_fields = [col for col in key_fields if col in existing_df.columns and col in new_data.columns]

            existing_combinations = existing_df.select(available_key_fields).unique()

            # Select new data columns that match what we can process
            new_select_cols = available_key_fields.copy()
            if "has_pricing" in new_data.columns:
                new_select_cols.append("has_pricing")

            new_combinations = new_data.select(new_select_cols).unique()

            # Anti-join to get only truly new combinations
            filtered_new = new_combinations.join(
                existing_combinations,
                on=available_key_fields,
                how="anti",
                join_nulls=True  # Allow nulls to match in join
            )

            new_count = filtered_new.height

            if new_count > 0:
                # Ensure filtered_new has same schema as existing_df before concat
                filtered_new = self._ensure_compatible_schema(filtered_new, existing_df)

                # Merge with existing data
                merged_df = pl.concat([existing_df, filtered_new], how="vertical")

                # Save to cache
                merged_df.write_parquet(self.cache_path)

                # Mark as modified
                self._mark_modified("add", new_count)

                self.logger.info(f"Added {new_count} new combinations to local cache")

            return new_count

        except Exception as e:
            self.logger.error(f"Failed to add records to cache: {e}")
            raise RuntimeError(f"Failed to add records to local cache: {e}") from e

    def flush_records(self, filters: FilterCriteria) -> FlushResult:
        """
        Remove records matching filter criteria from local cache.

        Args:
            filters: Criteria for selecting records to remove

        Returns:
            FlushResult with operation details
        """
        try:
            if not self.cache_path.exists():
                return FlushResult(
                    records_removed=0,
                    local_changes_made=False,
                    sync_reminder="No local cache file exists"
                )

            # Load current data
            df = pl.read_parquet(self.cache_path)

            # Apply filters to find records to remove
            records_to_remove = self._apply_filters(df, filters)
            removal_count = records_to_remove.height

            if removal_count == 0:
                return FlushResult(
                    records_removed=0,
                    local_changes_made=False,
                    sync_reminder="No records matched filter criteria"
                )

            # Remove matching records using anti-join
            key_fields = ["cloud_provider", "service", "usage_family",
                         "resource_type", "operation", "usage_type", "instance_type"]

            remaining_df = df.join(
                records_to_remove.select(key_fields).unique(),
                on=key_fields,
                how="anti",
                join_nulls=True  # Allow nulls to match in join
            )

            # Save updated cache
            remaining_df.write_parquet(self.cache_path)

            # Mark as modified
            self._mark_modified("flush", removal_count)

            self.logger.info(f"Removed {removal_count} records from local cache")

            return FlushResult(
                records_removed=removal_count,
                local_changes_made=True,
                sync_reminder="Run 'czbg data upload' to sync changes to S3"
            )

        except Exception as e:
            self.logger.error(f"Failed to flush records from cache: {e}")
            raise RuntimeError(f"Failed to flush records from local cache: {e}") from e

    def preview_flush(self, filters: FilterCriteria) -> FlushPreview:
        """
        Preview what records would be removed without actually removing them.

        Args:
            filters: Criteria for selecting records to preview

        Returns:
            FlushPreview with details about what would be removed
        """
        try:
            if not self.cache_path.exists():
                return FlushPreview(
                    records_to_remove=0,
                    affected_combinations=0,
                    sample_records=[]
                )

            df = pl.read_parquet(self.cache_path)
            records_to_remove = self._apply_filters(df, filters)

            removal_count = records_to_remove.height
            sample_limit = min(5, removal_count)
            sample_records = []

            if removal_count > 0:
                sample_df = records_to_remove.head(sample_limit)
                sample_records = sample_df.to_dicts()

            # Calculate affected combinations
            key_fields = ["cloud_provider", "service", "usage_family",
                         "resource_type", "operation", "usage_type", "instance_type"]
            affected_combinations = records_to_remove.select(key_fields).unique().height

            return FlushPreview(
                records_to_remove=removal_count,
                affected_combinations=affected_combinations,
                sample_records=sample_records
            )

        except Exception as e:
            self.logger.error(f"Failed to preview flush: {e}")
            return FlushPreview(
                records_to_remove=0,
                affected_combinations=0,
                sample_records=[]
            )

    def get_cache_info(self) -> Optional[CacheInfo]:
        """
        Get information about the local cache file.

        Returns:
            CacheInfo with file metadata or None if cache doesn't exist
        """
        try:
            if not self.cache_path.exists():
                return None

            # Get file stats
            stat = self.cache_path.stat()
            file_size_mb = stat.st_size / (1024 * 1024)
            modified_time = datetime.fromtimestamp(stat.st_mtime)

            # Read data for record count and checksum
            df = pl.read_parquet(self.cache_path)
            record_count = df.height

            # Calculate checksum
            checksum = self._calculate_file_checksum(self.cache_path)

            # Check if there are local changes
            has_local_changes = self._has_local_changes()

            return CacheInfo(
                file_path=self.cache_path,
                file_size_mb=file_size_mb,
                checksum=checksum,
                modified_time=modified_time,
                record_count=record_count,
                has_local_changes=has_local_changes
            )

        except Exception as e:
            self.logger.error(f"Failed to get cache info: {e}")
            return None

    def mark_synced(self, sync_version: str, sync_checksum: str):
        """
        Mark the local cache as synced with S3.

        Args:
            sync_version: Version ID synced with
            sync_checksum: Checksum of synced data
        """
        try:
            manifest = {
                "last_sync": datetime.utcnow().isoformat(),
                "sync_version": sync_version,
                "sync_checksum": sync_checksum,
                "has_local_changes": False
            }

            with open(self.local_manifest_path, 'w') as f:
                json.dump(manifest, f, indent=2)

            self.logger.debug(f"Marked cache as synced with version {sync_version}")

        except Exception as e:
            self.logger.error(f"Failed to update sync status: {e}")

    def _apply_filters(self, df: pl.DataFrame, filters: FilterCriteria) -> pl.DataFrame:
        """Apply filter criteria to find matching records."""
        filtered_df = df

        # Apply cloud provider filter (required)
        if filters.exact_match:
            filtered_df = filtered_df.filter(pl.col("cloud_provider") == filters.cloud_provider)
        else:
            filtered_df = filtered_df.filter(
                pl.col("cloud_provider").str.contains(filters.cloud_provider, literal=False)
            )

        # Apply optional filters
        if filters.service:
            if filters.exact_match:
                filtered_df = filtered_df.filter(pl.col("service") == filters.service)
            else:
                filtered_df = filtered_df.filter(
                    pl.col("service").str.contains(filters.service, literal=False)
                )

        if filters.usage_family:
            if filters.exact_match:
                filtered_df = filtered_df.filter(pl.col("usage_family") == filters.usage_family)
            else:
                filtered_df = filtered_df.filter(
                    pl.col("usage_family").str.contains(filters.usage_family, literal=False)
                )

        if filters.resource_type:
            if filters.exact_match:
                filtered_df = filtered_df.filter(pl.col("resource_type") == filters.resource_type)
            else:
                filtered_df = filtered_df.filter(
                    pl.col("resource_type").str.contains(filters.resource_type, literal=False)
                )

        if filters.operation:
            if filters.exact_match:
                filtered_df = filtered_df.filter(pl.col("operation") == filters.operation)
            else:
                filtered_df = filtered_df.filter(
                    pl.col("operation").str.contains(filters.operation, literal=False)
                )

        if filters.usage_type:
            if filters.exact_match:
                filtered_df = filtered_df.filter(pl.col("usage_type") == filters.usage_type)
            else:
                filtered_df = filtered_df.filter(
                    pl.col("usage_type").str.contains(filters.usage_type, literal=False)
                )

        return filtered_df

    def _mark_modified(self, operation: str, record_count: int):
        """Mark the cache as having local changes."""
        try:
            manifest = {
                "last_modified": datetime.utcnow().isoformat(),
                "last_operation": operation,
                "records_affected": record_count,
                "has_local_changes": True
            }

            # Merge with existing manifest
            if self.local_manifest_path.exists():
                with open(self.local_manifest_path, 'r') as f:
                    existing = json.load(f)
                existing.update(manifest)
                manifest = existing

            with open(self.local_manifest_path, 'w') as f:
                json.dump(manifest, f, indent=2)

        except Exception as e:
            self.logger.error(f"Failed to mark cache as modified: {e}")

    def _has_local_changes(self) -> bool:
        """Check if there are local changes since last sync."""
        try:
            if not self.local_manifest_path.exists():
                return True  # No sync record means changes exist

            with open(self.local_manifest_path, 'r') as f:
                manifest = json.load(f)

            return manifest.get("has_local_changes", True)

        except Exception as e:
            self.logger.error(f"Failed to check local changes: {e}")
            return True  # Assume changes exist on error

    def _get_standard_schema(self) -> dict:
        """Get the standard schema for reference data including all expected columns."""
        return {
            "cloud_provider": pl.Utf8,
            "service": pl.Utf8,
            "usage_family": pl.Utf8,
            "resource_type": pl.Utf8,
            "operation": pl.Utf8,
            "usage_type": pl.Utf8,
            "instance_type": pl.Utf8,  # Now included in standard schema
            "pricing_unit": pl.Utf8,
            "cost_per_unit": pl.Float64,
            "has_pricing": pl.Boolean
        }

    def _harmonize_schemas(self, existing_df: pl.DataFrame, new_df: pl.DataFrame) -> tuple[pl.DataFrame, pl.DataFrame]:
        """Harmonize schemas between existing and new data to prevent width mismatches."""
        standard_schema = self._get_standard_schema()

        # Add missing columns to existing_df with appropriate defaults
        for col, dtype in standard_schema.items():
            if col not in existing_df.columns:
                if dtype == pl.Boolean:
                    default_val = False
                elif dtype == pl.Float64:
                    default_val = 0.0
                else:
                    default_val = ""
                existing_df = existing_df.with_columns([pl.lit(default_val).cast(dtype).alias(col)])

        # Add missing columns to new_df with appropriate defaults
        for col, dtype in standard_schema.items():
            if col not in new_df.columns:
                if dtype == pl.Boolean:
                    default_val = False
                elif dtype == pl.Float64:
                    default_val = 0.0
                else:
                    default_val = ""
                new_df = new_df.with_columns([pl.lit(default_val).cast(dtype).alias(col)])

        # Ensure column order matches standard schema
        column_order = list(standard_schema.keys())
        existing_cols = [col for col in column_order if col in existing_df.columns]
        new_cols = [col for col in column_order if col in new_df.columns]

        existing_df = existing_df.select(existing_cols)
        new_df = new_df.select(new_cols)

        return existing_df, new_df

    def _ensure_compatible_schema(self, new_df: pl.DataFrame, target_df: pl.DataFrame) -> pl.DataFrame:
        """Ensure new_df has compatible schema with target_df for concatenation."""
        # Add missing columns from target_df to new_df
        for col in target_df.columns:
            if col not in new_df.columns:
                # Get the data type from target_df
                target_dtype = target_df.select(col).dtypes[0]

                if target_dtype == pl.Boolean:
                    default_val = False
                elif target_dtype == pl.Float64:
                    default_val = 0.0
                else:
                    default_val = ""

                new_df = new_df.with_columns([pl.lit(default_val).cast(target_dtype).alias(col)])

        # Select columns in same order as target_df
        return new_df.select(target_df.columns)

    def deduplicate_combinations(self, df: pl.DataFrame) -> tuple[pl.DataFrame, int]:
        """
        Remove duplicate combinations, prioritizing records with usage data.

        Args:
            df: DataFrame with combination and usage data

        Returns:
            Tuple of (deduplicated DataFrame, count of duplicates removed)
        """
        if df.is_empty():
            return df, 0

        original_count = df.height
        deduplicated_df = self._apply_deduplication_logic(df)
        duplicates_removed = original_count - deduplicated_df.height

        return deduplicated_df, duplicates_removed

    def _apply_deduplication_logic(self, df: pl.DataFrame) -> pl.DataFrame:
        """Apply deduplication logic prioritizing records with usage data."""
        combination_fields = [
            "cloud_provider", "service", "usage_family", "resource_type",
            "operation", "usage_type", "instance_type"
        ]

        # Filter to existing combination fields only
        existing_fields = [field for field in combination_fields if field in df.columns]
        if not existing_fields:
            return df

        # Add usage_amount column if missing (for prioritization)
        if "usage_amount" not in df.columns:
            df = df.with_columns([pl.lit(0.0).alias("usage_amount")])

        # Sort by usage_amount descending, then keep first of each group
        return (
            df.sort(["usage_amount"], descending=[True])
            .unique(subset=existing_fields, keep="first")
        )

    def _get_combination_fields(self) -> list[str]:
        """Get the 7 key fields that define a unique combination."""
        return [
            "cloud_provider", "service", "usage_family", "resource_type",
            "operation", "usage_type", "instance_type"
        ]

    def _clean_and_deduplicate_cache(self, cache_df: pl.DataFrame) -> dict:
        """Clean resource_type prefixes and remove duplicates from cache data."""
        if cache_df.is_empty():
            return {
                "cleaned_df": cache_df,
                "resource_types_cleaned": 0,
                "duplicates_removed": 0
            }

        # Count resource types that need cleaning
        resource_types_cleaned = self._count_cache_resource_types_needing_cleaning(cache_df)

        # Clean resource_type prefixes
        if "resource_type" in cache_df.columns:
            cache_df = cache_df.with_columns([
                pl.col("resource_type")
                .str.replace(r"^[^:]+:\s*", "")  # Remove any "prefix: " pattern
                .str.strip_chars()
                .alias("resource_type")
            ])

        # Deduplicate combinations
        deduplicated_df, duplicates_removed = self.deduplicate_combinations(cache_df)

        return {
            "cleaned_df": deduplicated_df,
            "resource_types_cleaned": resource_types_cleaned,
            "duplicates_removed": duplicates_removed
        }

    def _count_cache_resource_types_needing_cleaning(self, df: pl.DataFrame) -> int:
        """Count resource_type values in cache that need prefix cleaning."""
        if "resource_type" not in df.columns or df.is_empty():
            return 0

        return df.filter(
            pl.col("resource_type").is_not_null() &
            pl.col("resource_type").str.contains(r"^[^:]+:\s*")
        ).height

    def merge_dual_operations(self, combinations: pl.DataFrame, pricing: pl.DataFrame) -> dict:
        """
        Merge results from both operations (combinations + pricing) into cache.

        Always updates pricing for existing combinations when available.

        Args:
            combinations: DataFrame with unique service combinations (Operation 1)
            pricing: DataFrame with cost-per-unit calculations (Operation 2)

        Returns:
            Dictionary with detailed operation results
        """
        try:
            # If no pricing data, fall back to combinations-only
            if pricing.is_empty():
                new_combinations = self.add_records(combinations)
                return {
                    "new_combinations": new_combinations,
                    "pricing_updates": 0,
                    "total_combinations_processed": combinations.height
                }

            # Execute dual operations: new combinations + pricing updates
            result = self._execute_dual_operations(combinations, pricing)
            return result

        except Exception as e:
            self.logger.error(f"Failed to merge dual operations: {e}")
            raise RuntimeError(f"Failed to merge dual operations: {e}") from e

    def _execute_dual_operations(self,
                               combinations: pl.DataFrame,
                               pricing: pl.DataFrame) -> dict:
        """Execute both operations: track new combinations and update all pricing."""
        # Load existing data
        if self.cache_path.exists():
            existing_df = pl.read_parquet(self.cache_path)
        else:
            existing_df = pl.DataFrame(schema=self._get_standard_schema())

        # Clean and deduplicate cache data (always)
        cache_stats = self._clean_and_deduplicate_cache(existing_df)
        existing_df = cache_stats["cleaned_df"]

        # Harmonize schemas
        existing_df, combinations = self._harmonize_schemas(existing_df, combinations)
        existing_df, pricing = self._harmonize_schemas(existing_df, pricing)

        # Step 1: Add new combinations (Operation 1 - unchanged)
        new_combinations = self._add_new_combinations_only(existing_df, combinations)

        # Step 2: Update ALL pricing data (Operation 1 enhancement)
        pricing_updates = self._update_all_pricing(existing_df, pricing)

        # Step 3: Merge and save final result
        final_df = self._create_final_dataset(existing_df, new_combinations, pricing)

        # Save to cache
        final_df.write_parquet(self.cache_path)

        # Mark as modified if any changes occurred
        total_changes = new_combinations + pricing_updates
        if total_changes > 0:
            self._mark_modified("dual_operations", total_changes)

        self.logger.info(
            f"Dual operations completed: {new_combinations} new combinations, "
            f"{pricing_updates} pricing updates"
        )

        # Return combined statistics
        return {
            "new_combinations": new_combinations,
            "pricing_updates": pricing_updates,
            "total_combinations_processed": combinations.height,
            "cache_resource_types_cleaned": cache_stats["resource_types_cleaned"],
            "cache_duplicates_removed": cache_stats["duplicates_removed"]
        }

    def _add_new_combinations_only(self,
                                 existing_df: pl.DataFrame,
                                 combinations: pl.DataFrame) -> int:
        """Add only genuinely new combinations, return count."""
        key_fields = [
            "cloud_provider", "service", "usage_family",
            "resource_type", "operation", "usage_type", "instance_type"
        ]

        available_key_fields = [
            col for col in key_fields
            if col in existing_df.columns and col in combinations.columns
        ]

        if not available_key_fields:
            return 0

        existing_combinations = existing_df.select(available_key_fields).unique()
        new_combinations = combinations.join(
            existing_combinations,
            on=available_key_fields,
            how="anti",
            join_nulls=True  # Allow nulls to match in join
        )

        return new_combinations.height

    def _update_all_pricing(self,
                          existing_df: pl.DataFrame,
                          pricing: pl.DataFrame) -> int:
        """Update pricing for ALL combinations that have pricing data."""
        if pricing.is_empty():
            return 0

        key_fields = [
            "cloud_provider", "service", "usage_family",
            "resource_type", "operation", "usage_type", "instance_type"
        ]

        available_key_fields = [
            col for col in key_fields
            if col in existing_df.columns and col in pricing.columns
        ]

        if not available_key_fields:
            return 0

        # Find existing combinations that have pricing updates
        existing_combinations = existing_df.select(available_key_fields).unique()
        pricing_matches = pricing.join(
            existing_combinations,
            on=available_key_fields,
            how="inner",
            join_nulls=True  # Allow nulls to match in join
        )

        return pricing_matches.height

    def _create_final_dataset(self,
                            existing_df: pl.DataFrame,
                            new_combination_count: int,
                            pricing: pl.DataFrame) -> pl.DataFrame:
        """Create final dataset with new combinations and updated pricing."""
        # Start with existing data
        final_df = existing_df.clone()

        # If we have pricing data, update/merge it
        if not pricing.is_empty():
            key_fields = [
                "cloud_provider", "service", "usage_family",
                "resource_type", "operation", "usage_type", "instance_type"
            ]

            available_keys = [
                col for col in key_fields
                if col in final_df.columns and col in pricing.columns
            ]

            if available_keys:
                # Update existing records with new pricing
                # Use left join to preserve all existing data, update pricing where available
                final_df = final_df.join(
                    pricing.select(available_keys + ["cost_per_unit", "pricing_unit"]),
                    on=available_keys,
                    how="left",
                    suffix="_new",
                    join_nulls=True  # Allow nulls to match in join
                )

                # Coalesce new pricing with existing, preferring new data
                if "cost_per_unit_new" in final_df.columns:
                    final_df = final_df.with_columns([
                        pl.coalesce(["cost_per_unit_new", "cost_per_unit"]).alias("cost_per_unit")
                    ]).drop(["cost_per_unit_new"])

                if "pricing_unit_new" in final_df.columns:
                    final_df = final_df.with_columns([
                        pl.coalesce(["pricing_unit_new", "pricing_unit"]).alias("pricing_unit")
                    ]).drop(["pricing_unit_new"])

                # Update has_pricing flag
                final_df = final_df.with_columns([
                    (pl.col("cost_per_unit").is_not_null() &
                     (pl.col("cost_per_unit") > 0)).alias("has_pricing")
                ])

                # Add any truly new combinations from pricing that weren't in existing data
                new_from_pricing = pricing.join(
                    final_df.select(available_keys),
                    on=available_keys,
                    how="anti",
                    join_nulls=True  # Allow nulls to match in join
                )

                if not new_from_pricing.is_empty():
                    # Ensure compatible schema for new pricing records
                    new_from_pricing = self._ensure_compatible_schema(new_from_pricing, final_df)
                    final_df = pl.concat([final_df, new_from_pricing], how="vertical")

        return final_df

    def _calculate_file_checksum(self, file_path: Path) -> str:
        """Calculate MD5 checksum of a file."""
        try:
            md5_hash = hashlib.md5()
            with open(file_path, "rb") as f:
                for chunk in iter(lambda: f.read(4096), b""):
                    md5_hash.update(chunk)
            return md5_hash.hexdigest()
        except Exception as e:
            self.logger.error(f"Failed to calculate checksum: {e}")
            return ""