#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0

"""Analyze chunk hash files generated by the file_hash strategy.

This script reads chunk hash files and provides statistics about chunk reuse patterns.
"""

# Standard
from collections import Counter, defaultdict
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional, Set
import argparse
import json
import sys


def load_chunk_hashes(input_dir: Path) -> List[Dict]:
    """Load all chunk hash records from JSONL files.

    Args:
        input_dir: Directory containing chunk hash JSONL files

    Returns:
        List of chunk hash records
    """
    records: List[Dict] = []
    hash_files = sorted(input_dir.glob("chunk_hashes_*.jsonl"))

    if not hash_files:
        print(f"No chunk hash files found in {input_dir}")
        return records

    print(f"Found {len(hash_files)} chunk hash files")

    for hash_file in hash_files:
        print(f"Reading {hash_file.name}...")
        with open(hash_file) as f:
            for line in f:
                try:
                    data = json.loads(line)
                    records.append(data)
                except json.JSONDecodeError as e:
                    print(f"Warning: Failed to parse line in {hash_file}: {e}")
                    continue

    return records


def analyze_chunk_hashes(records: List[Dict], top_n: int = 10) -> Dict:
    """Analyze chunk hash records and compute statistics.

    Args:
        records: List of chunk hash records
        top_n: Number of top reused chunks to include

    Returns:
        Dictionary containing analysis results
    """
    all_hashes: List[str] = []
    unique_hashes: Set[str] = set()
    lookup_ids: Set[str] = set()
    hash_frequency: Counter = Counter()
    request_stats: Dict[str, Dict] = {}

    for record in records:
        lookup_id = record.get("lookup_id", "unknown")
        chunk_hashes = record.get("chunk_hashes", [])

        lookup_ids.add(lookup_id)
        all_hashes.extend(chunk_hashes)
        unique_hashes.update(chunk_hashes)
        hash_frequency.update(chunk_hashes)

        if lookup_id not in request_stats:
            request_stats[lookup_id] = {
                "chunk_count": len(chunk_hashes),
                "timestamp": record.get("timestamp", 0),
            }

    total_chunks = len(all_hashes)
    unique_chunks = len(unique_hashes)
    duplicate_chunks = total_chunks - unique_chunks
    reuse_rate = duplicate_chunks / total_chunks if total_chunks > 0 else 0.0

    # Find most frequently reused chunks
    top_reused = hash_frequency.most_common(top_n)

    # Per-request statistics
    chunk_counts = [stats["chunk_count"] for stats in request_stats.values()]
    avg_chunks = sum(chunk_counts) / len(chunk_counts) if chunk_counts else 0
    max_chunks = max(chunk_counts) if chunk_counts else 0
    min_chunks = min(chunk_counts) if chunk_counts else 0

    # Frequency distribution
    reuse_counts = list(hash_frequency.values())
    single_use = sum(1 for c in reuse_counts if c == 1)
    multi_use = len(reuse_counts) - single_use

    return {
        "total_records": len(records),
        "total_chunks": total_chunks,
        "unique_chunks": unique_chunks,
        "duplicate_chunks": duplicate_chunks,
        "reuse_rate": reuse_rate,
        "unique_lookup_ids": len(lookup_ids),
        "top_reused_chunks": top_reused,
        "request_stats": {
            "total_requests": len(request_stats),
            "avg_chunks_per_request": avg_chunks,
            "max_chunks_per_request": max_chunks,
            "min_chunks_per_request": min_chunks,
        },
        "frequency_distribution": {
            "single_use_chunks": single_use,
            "multi_use_chunks": multi_use,
            "single_use_percentage": single_use / len(reuse_counts)
            if reuse_counts
            else 0,
            "multi_use_percentage": multi_use / len(reuse_counts)
            if reuse_counts
            else 0,
        },
    }


def print_analysis(analysis: Dict, verbose: bool = False) -> None:
    """Print analysis results in a readable format.

    Args:
        analysis: Analysis results dictionary
        verbose: Whether to print detailed statistics
    """
    print("\n" + "=" * 60)
    print("Chunk Hash Analysis Results")
    print("=" * 60)

    print("\nBasic Statistics:")
    print(f"  Total records:        {analysis['total_records']:,}")
    print(f"  Unique lookup IDs:    {analysis['unique_lookup_ids']:,}")
    print(f"  Total chunks:         {analysis['total_chunks']:,}")
    print(f"  Unique chunks:        {analysis['unique_chunks']:,}")
    print(f"  Duplicate chunks:     {analysis['duplicate_chunks']:,}")
    print(f"  Reuse rate:           {analysis['reuse_rate']:.2%}")

    if verbose:
        req_stats = analysis["request_stats"]
        print("\nPer-Request Statistics:")
        print(f"  Total requests:       {req_stats['total_requests']:,}")
        print(f"  Avg chunks/request:   {req_stats['avg_chunks_per_request']:.2f}")
        print(f"  Max chunks/request:   {req_stats['max_chunks_per_request']:,}")
        print(f"  Min chunks/request:   {req_stats['min_chunks_per_request']:,}")

        freq_dist = analysis["frequency_distribution"]
        print("\nFrequency Distribution:")
        single_pct = freq_dist["single_use_percentage"]
        multi_pct = freq_dist["multi_use_percentage"]
        print(
            f"  Single-use chunks:    "
            f"{freq_dist['single_use_chunks']:,} ({single_pct:.2%})"
        )
        print(
            f"  Multi-use chunks:     "
            f"{freq_dist['multi_use_chunks']:,} ({multi_pct:.2%})"
        )

    top_n = len(analysis["top_reused_chunks"])
    print(f"\nTop {top_n} Most Reused Chunks:")
    for i, (chunk_hash, count) in enumerate(analysis["top_reused_chunks"], 1):
        print(f"  {i:2d}. {chunk_hash[:16]}... (reused {count:,} times)")

    print("\n" + "=" * 60)


def analyze_time_series(records: List[Dict]) -> Dict:
    """Analyze chunk distribution over time.

    Args:
        records: List of chunk hash records

    Returns:
        Dictionary containing time-series statistics
    """
    hourly_stats: Dict = defaultdict(lambda: {"total": 0, "unique": set()})

    for record in records:
        timestamp = record.get("timestamp", 0)
        if timestamp > 0:
            hour = datetime.fromtimestamp(timestamp).strftime("%Y-%m-%d %H:00")
            chunk_hashes = record.get("chunk_hashes", [])
            hourly_stats[hour]["total"] += len(chunk_hashes)
            hourly_stats[hour]["unique"].update(chunk_hashes)

    # Convert sets to counts for JSON serialization
    time_series = {}
    for hour, stats in sorted(hourly_stats.items()):
        total = stats["total"]
        unique = len(stats["unique"])
        reuse_rate = (total - unique) / total if total > 0 else 0.0
        time_series[hour] = {
            "total_chunks": total,
            "unique_chunks": unique,
            "reuse_rate": reuse_rate,
        }

    return time_series


def estimate_memory(
    unique_chunks: int, chunk_size: int = 256, bytes_per_token: int = 2
) -> Dict:
    """Estimate memory requirements for caching.

    Args:
        unique_chunks: Number of unique chunks
        chunk_size: Number of tokens per chunk
        bytes_per_token: Bytes per token

    Returns:
        Dictionary containing memory estimates
    """
    bytes_per_chunk = chunk_size * bytes_per_token
    total_memory_bytes = unique_chunks * bytes_per_chunk
    total_memory_mb = total_memory_bytes / (1024 * 1024)
    total_memory_gb = total_memory_mb / 1024

    return {
        "unique_chunks": unique_chunks,
        "chunk_size_tokens": chunk_size,
        "bytes_per_token": bytes_per_token,
        "bytes_per_chunk": bytes_per_chunk,
        "total_memory_bytes": total_memory_bytes,
        "total_memory_mb": total_memory_mb,
        "total_memory_gb": total_memory_gb,
    }


def export_results(
    analysis: Dict,
    output_file: Path,
    include_time_series: bool = False,
    records: Optional[List[Dict]] = None,
) -> None:
    """Export analysis results to a JSON file.

    Args:
        analysis: Analysis results dictionary
        output_file: Output file path
        include_time_series: Whether to include time-series analysis
        records: Original records (needed for time-series analysis)
    """
    # Convert Counter objects to lists for JSON serialization
    export_data = analysis.copy()
    export_data["top_reused_chunks"] = [
        {"hash": h, "count": c} for h, c in analysis["top_reused_chunks"]
    ]

    # Add memory estimation
    export_data["memory_estimation"] = estimate_memory(analysis["unique_chunks"])

    # Add time-series analysis if requested
    if include_time_series and records:
        export_data["time_series"] = analyze_time_series(records)

    with open(output_file, "w") as f:
        json.dump(export_data, f, indent=2)

    print(f"\nResults exported to {output_file}")


def print_time_series(time_series: Dict) -> None:
    """Print time-series analysis results.

    Args:
        time_series: Time-series statistics dictionary
    """
    if not time_series:
        print("\nNo time-series data available")
        return

    print("\n" + "=" * 70)
    print("Time-Series Analysis")
    print("=" * 70)
    header = (
        f"\n{'Hour':<20} | {'Total Chunks':>12} | "
        f"{'Unique Chunks':>13} | {'Reuse Rate':>10}"
    )
    print(header)
    print("-" * 70)

    for hour, stats in sorted(time_series.items()):
        row = (
            f"{hour:<20} | {stats['total_chunks']:>12,} | "
            f"{stats['unique_chunks']:>13,} | {stats['reuse_rate']:>9.2%}"
        )
        print(row)

    print("=" * 70)


def print_memory_estimation(memory_est: Dict) -> None:
    """Print memory estimation results.

    Args:
        memory_est: Memory estimation dictionary
    """
    print("\n" + "=" * 60)
    print("Memory Estimation")
    print("=" * 60)
    print(f"\n  Unique chunks:        {memory_est['unique_chunks']:,}")
    print(f"  Chunk size:           {memory_est['chunk_size_tokens']} tokens")
    print(f"  Bytes per token:      {memory_est['bytes_per_token']} bytes")
    print(f"  Bytes per chunk:      {memory_est['bytes_per_chunk']:,} bytes")
    mem_mb = memory_est["total_memory_mb"]
    mem_gb = memory_est["total_memory_gb"]
    print(f"  Total memory:         {mem_mb:.2f} MB ({mem_gb:.2f} GB)")
    print("=" * 60)


def main():
    parser = argparse.ArgumentParser(
        description="Analyze chunk hash files from LMCache file_hash strategy"
    )
    parser.add_argument(
        "--input-dir",
        type=Path,
        default=Path("./chunk_hashes"),
        help=("Directory containing chunk hash JSONL files (default: ./chunk_hashes)"),
    )
    parser.add_argument(
        "--output",
        type=Path,
        help="Output JSON file for analysis results (optional)",
    )
    parser.add_argument(
        "--top-n",
        type=int,
        default=10,
        help="Number of top reused chunks to display (default: 10)",
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help=(
            "Show detailed statistics including per-request and frequency distribution"
        ),
    )
    parser.add_argument(
        "--time-series",
        action="store_true",
        help="Include time-series analysis (hourly breakdown)",
    )
    parser.add_argument(
        "--memory-estimation",
        action="store_true",
        help="Show memory estimation for caching unique chunks",
    )
    parser.add_argument(
        "--chunk-size",
        type=int,
        default=256,
        help=("Chunk size in tokens for memory estimation (default: 256)"),
    )
    parser.add_argument(
        "--bytes-per-token",
        type=int,
        default=2,
        help="Bytes per token for memory estimation (default: 2)",
    )

    args = parser.parse_args()

    if not args.input_dir.exists():
        print(f"Error: Input directory {args.input_dir} does not exist")
        return 1

    if not args.input_dir.is_dir():
        print(f"Error: {args.input_dir} is not a directory")
        return 1

    # Load and analyze chunk hashes
    print(f"Loading chunk hashes from {args.input_dir}...")
    records = load_chunk_hashes(args.input_dir)

    if not records:
        print("No records found to analyze")
        return 1

    print(f"Analyzing {len(records)} records...")
    analysis = analyze_chunk_hashes(records, top_n=args.top_n)

    # Print results
    print_analysis(analysis, verbose=args.verbose)

    # Time-series analysis
    if args.time_series:
        time_series = analyze_time_series(records)
        print_time_series(time_series)

    # Memory estimation
    if args.memory_estimation:
        memory_est = estimate_memory(
            analysis["unique_chunks"],
            chunk_size=args.chunk_size,
            bytes_per_token=args.bytes_per_token,
        )
        print_memory_estimation(memory_est)

    # Export results if requested
    if args.output:
        export_results(
            analysis,
            args.output,
            include_time_series=args.time_series,
            records=records if args.time_series else None,
        )

    return 0


if __name__ == "__main__":
    sys.exit(main())
