from __future__ import annotations

import logging
import time
from collections import Counter
from collections import OrderedDict
from dataclasses import dataclass
from dataclasses import field
from typing import Any
from typing import Dict

NULL_FLAG = -(1 << 63)
MIN_K_HASHES = 32
HISTOGRAM_BINS = 32


@dataclass
class DataFile:
    file_path: str
    file_format: str = "PARQUET"
    record_count: int = 0
    file_size_in_bytes: int = 0
    partition: Dict[str, object] = field(default_factory=dict)
    lower_bounds: Dict[int, bytes] | None = None
    upper_bounds: Dict[int, bytes] | None = None


@dataclass
class ManifestEntry:
    snapshot_id: int
    data_file: DataFile
    status: str = "added"  # 'added' | 'deleted'


@dataclass
class ParquetManifestEntry:
    """Represents a single entry in a Parquet manifest with statistics."""

    file_path: str
    file_format: str
    record_count: int
    file_size_in_bytes: int
    uncompressed_size_in_bytes: int
    column_uncompressed_sizes_in_bytes: list[int]
    null_counts: list[int]
    min_k_hashes: list[list[int]]
    histogram_counts: list[list[int]]
    histogram_bins: int
    min_values: list
    max_values: list
    min_values_display: list
    max_values_display: list

    def to_dict(self) -> dict:
        return {
            "file_path": self.file_path,
            "file_format": self.file_format,
            "record_count": self.record_count,
            "file_size_in_bytes": self.file_size_in_bytes,
            "uncompressed_size_in_bytes": self.uncompressed_size_in_bytes,
            "column_uncompressed_sizes_in_bytes": self.column_uncompressed_sizes_in_bytes,
            "null_counts": self.null_counts,
            "min_k_hashes": self.min_k_hashes,
            "histogram_counts": self.histogram_counts,
            "histogram_bins": self.histogram_bins,
            "min_values": self.min_values,
            "max_values": self.max_values,
            "min_values_display": self.min_values_display,
            "max_values_display": self.max_values_display,
        }


logger = logging.getLogger(__name__)
_manifest_metrics = Counter()

# Parsed-manifest cache (LRU): store parsed Python representation (list[dict])
# to avoid repeated pyarrow parsing and expensive to_pylist() conversions.
# Entries are "frozen" for memory efficiency (inner lists -> tuples).
PARSED_MANIFEST_CACHE_SIZE: int = 32
_parsed_manifest_cache: "OrderedDict[str, list]" = OrderedDict()


def _freeze_for_cache(value):
    """Recursively freeze lists to tuples and convert byte-like to bytes.

    Keeps top-level entries as dicts (callers expect Mapping access) but
    replaces inner mutable lists with tuples to reduce memory overhead and
    prevent accidental mutation of cached data.
    """
    # Primitive/bytes
    if isinstance(value, (bytes, bytearray, memoryview)):
        return bytes(value)
    if isinstance(value, (str, int, float, bool)) or value is None:
        return value

    # Lists -> tuples (recursive)
    if isinstance(value, list):
        return tuple(_freeze_for_cache(v) for v in value)

    # Dicts -> keep as dict but freeze values
    if isinstance(value, dict):
        return {k: _freeze_for_cache(v) for k, v in value.items()}

    # Fallback: return as-is
    return value


def get_parsed_manifest(io, manifest_path: str) -> list:
    """Return a cached Python representation (list[dict]) of the Parquet manifest.

    - Uses an in-memory LRU cache keyed by `manifest_path`.
    - Cached entries are frozen (inner lists -> tuples) to reduce memory.
    - Callers MUST treat returned lists/dicts as read-only.
    """

    if not manifest_path:
        return []

    # Fast path: cache hit
    if manifest_path in _parsed_manifest_cache:
        _parsed_manifest_cache.move_to_end(manifest_path)
        _manifest_metrics["parsed_cache_hits"] += 1
        return _parsed_manifest_cache[manifest_path]

    # Miss: read bytes -> parse -> freeze -> cache
    import pyarrow as pa
    import pyarrow.parquet as pq

    inp = io.new_input(manifest_path)
    try:
        with inp.open() as f:
            data = f.read()
    except FileNotFoundError:
        # keep behavior consistent with callers
        raise

    if not data:
        _manifest_metrics["parsed_cache_misses"] += 1
        _parsed_manifest_cache[manifest_path] = []
        # Evict oldest if needed
        if len(_parsed_manifest_cache) > PARSED_MANIFEST_CACHE_SIZE:
            _parsed_manifest_cache.popitem(last=False)
        return []

    table = pq.read_table(pa.BufferReader(data))
    rows = table.to_pylist()

    # Freeze inner mutable structures to tuples where possible
    frozen_rows: list = []
    for r in rows:
        if isinstance(r, dict):
            fr = {k: _freeze_for_cache(v) for k, v in r.items()}
            frozen_rows.append(fr)
        else:
            frozen_rows.append(r)

    _parsed_manifest_cache[manifest_path] = frozen_rows
    _manifest_metrics["parsed_cache_misses"] += 1

    # Evict oldest if cache exceeds size
    if len(_parsed_manifest_cache) > PARSED_MANIFEST_CACHE_SIZE:
        _parsed_manifest_cache.popitem(last=False)

    return _parsed_manifest_cache[manifest_path]


def invalidate_parsed_manifest(manifest_path: str) -> None:
    """Remove a manifest from the parsed-manifest cache (if present)."""
    _parsed_manifest_cache.pop(manifest_path, None)


def clear_parsed_manifest_cache() -> None:
    """Clear the entire parsed-manifest cache (tests / admin use)."""
    _parsed_manifest_cache.clear()


def _compute_stats_for_arrow_column(col, field_type, file_path: str):
    """Compute statistics for a single PyArrow column (Array or ChunkedArray).

    Returns a tuple: (col_min_k, col_hist, col_min, col_max, min_display, max_display, null_count)
    """
    import heapq

    import opteryx.draken as draken  # type: ignore
    import pyarrow as pa

    # Ensure single contiguous array when possible
    if hasattr(col, "combine_chunks"):
        try:
            col = col.combine_chunks()
        except Exception:
            # leave as-is
            pass

    # Record compress/hash usage
    _manifest_metrics["hash_calls"] += 1
    _manifest_metrics["compress_calls"] += 1

    col_py = None
    try:
        vec = draken.Vector.from_arrow(col)
    except Exception:  # pragma: no cover - be robust
        raise

    hashes = set(vec.hash())

    # Decide whether to compute min-k/histogram for this column
    compute_min_k = False
    if (
        pa.types.is_integer(field_type)
        or pa.types.is_floating(field_type)
        or pa.types.is_decimal(field_type)
    ):
        compute_min_k = True
    elif (
        pa.types.is_timestamp(field_type)
        or pa.types.is_date(field_type)
        or pa.types.is_time(field_type)
    ):
        compute_min_k = True
    elif (
        pa.types.is_string(field_type)
        or pa.types.is_large_string(field_type)
        or pa.types.is_binary(field_type)
        or pa.types.is_large_binary(field_type)
    ):
        # For strings/binary we may need pylist for display
        try:
            col_py = col.to_pylist()
        except Exception:
            col_py = None
        compute_min_k = True

    if compute_min_k:
        smallest = heapq.nsmallest(MIN_K_HASHES, hashes)
        col_min_k = sorted(smallest)
    else:
        col_min_k = []

    import pyarrow as pa  # local import for types

    compute_hist = compute_min_k
    if pa.types.is_boolean(field_type):
        compute_hist = True

    # Use draken.compress() to get canonical int64 per value
    compressed = list(vec.compress())
    null_count = sum(1 for m in compressed if m == NULL_FLAG)

    non_nulls_compressed = [m for m in compressed if m != NULL_FLAG]
    if non_nulls_compressed:
        vmin = min(non_nulls_compressed)
        vmax = max(non_nulls_compressed)
        col_min = int(vmin)
        col_max = int(vmax)
        if compute_hist:
            # Special-case boolean histograms
            if pa.types.is_boolean(field_type):
                try:
                    if col_py is None:
                        try:
                            col_py = col.to_pylist()
                        except Exception:
                            col_py = None
                    if col_py is not None:
                        non_nulls_bool = [v for v in col_py if v is not None]
                        false_count = sum(1 for v in non_nulls_bool if v is False)
                        true_count = sum(1 for v in non_nulls_bool if v is True)
                    else:
                        # Fallback: infer from compressed mapping (assume 0/1)
                        false_count = sum(1 for m in non_nulls_compressed if m == 0)
                        true_count = sum(1 for m in non_nulls_compressed if m != 0)
                except Exception:
                    false_count = 0
                    true_count = 0

                col_hist = [int(true_count), int(false_count)]
            else:
                if vmin == vmax:
                    col_hist = []
                else:
                    col_hist = [0] * HISTOGRAM_BINS
                    span = float(vmax - vmin)
                    for m in non_nulls_compressed:
                        b = int(((float(m) - float(vmin)) / span) * (HISTOGRAM_BINS - 1))
                        if b < 0:
                            b = 0
                        if b >= HISTOGRAM_BINS:
                            b = HISTOGRAM_BINS - 1
                        col_hist[b] += 1
        else:
            col_hist = []
    else:
        # no non-null values
        col_min = NULL_FLAG
        col_max = NULL_FLAG
        col_hist = []

    # display values
    try:
        if pa.types.is_string(field_type) or pa.types.is_large_string(field_type):
            if col_py is None:
                try:
                    col_py = col.to_pylist()
                except Exception:
                    col_py = None
            if col_py is not None:
                non_nulls_str = [x for x in col_py if x is not None]
                if non_nulls_str:
                    min_value = min(non_nulls_str)
                    max_value = max(non_nulls_str)
                    if len(min_value) > 16:
                        min_value = min_value[:16] + "..."
                    if len(max_value) > 16:
                        max_value = max_value[:16] + "..."
                    min_display = min_value
                    max_display = max_value
                else:
                    min_display = None
                    max_display = None
            else:
                min_display = None
                max_display = None
        elif pa.types.is_binary(field_type) or pa.types.is_large_binary(field_type):
            if col_py is None:
                try:
                    col_py = col.to_pylist()
                except Exception:
                    col_py = None
            if col_py is not None:
                non_nulls = [x for x in col_py if x is not None]
                if non_nulls:
                    min_value = min(non_nulls)
                    max_value = max(non_nulls)
                    if len(min_value) > 16:
                        min_value = min_value[:16] + "..."
                    if len(max_value) > 16:
                        max_value = max_value[:16] + "..."
                    if any(ord(b) < 32 or ord(b) > 126 for b in min_value):
                        min_value = min_value.hex()
                        min_value = min_value[:16] + "..."
                    if any(ord(b) < 32 or ord(b) > 126 for b in max_value):
                        max_value = max_value.hex()
                        max_value = max_value[:16] + "..."
                    min_display = min_value
                    max_display = max_value
                else:
                    min_display = None
                    max_display = None
            else:
                min_display = None
                max_display = None
        else:
            # Handle dates and timestamps with special formatting
            if (
                pa.types.is_date(field_type)
                or pa.types.is_timestamp(field_type)
                or pa.types.is_time(field_type)
            ):
                if col_py is None:
                    try:
                        col_py = col.to_pylist()
                    except Exception:
                        col_py = None
                if col_py is not None:
                    non_nulls = [x for x in col_py if x is not None]
                    if non_nulls:
                        min_val = min(non_nulls)
                        max_val = max(non_nulls)
                        # Convert to ISO format strings for proper display
                        try:
                            min_display = (
                                min_val.isoformat()
                                if hasattr(min_val, "isoformat")
                                else str(min_val)
                            )
                            max_display = (
                                max_val.isoformat()
                                if hasattr(max_val, "isoformat")
                                else str(max_val)
                            )
                        except Exception:
                            min_display = str(min_val)
                            max_display = str(max_val)
                    else:
                        min_display = None
                        max_display = None
                else:
                    min_display = None
                    max_display = None
            else:
                # For other types (numbers, etc.)
                if col_py is None:
                    try:
                        col_py = col.to_pylist()
                    except Exception:
                        col_py = None
                if col_py is not None:
                    non_nulls = [x for x in col_py if x is not None]
                    if non_nulls:
                        min_display = min(non_nulls)
                        max_display = max(non_nulls)
                    else:
                        min_display = None
                        max_display = None
                else:
                    min_display = None
                    max_display = None
    except Exception:
        min_display = None
        max_display = None

    return (
        col_min_k,
        col_hist,
        int(col_min),
        int(col_max),
        min_display,
        max_display,
        int(null_count),
    )


def build_parquet_manifest_entry_from_bytes(
    data_bytes: bytes,
    file_path: str,
    file_size_in_bytes: int | None = None,
    orig_table: Any | None = None,
) -> ParquetManifestEntry:
    """Build a manifest entry by reading a parquet file as bytes and scanning column-by-column.

    This reads the compressed file once and materializes one full column at a time
    (combine_chunks) which keeps peak memory low while letting per-column
    stat calculation (draken) operate on contiguous arrays.
    """
    import pyarrow as pa
    import pyarrow.parquet as pq

    t_start = time.perf_counter()
    _manifest_metrics["files_read"] += 1
    _manifest_metrics["bytes_read"] += len(data_bytes)
    data_bytes_len = len(data_bytes)

    buf = pa.BufferReader(data_bytes)
    pf = pq.ParquetFile(buf)
    meta = pf.metadata

    # Try to read rugo metadata early so we can compute sizes without
    # materializing the table later. This is zero-copy and fast.
    try:
        from opteryx.rugo.parquet import read_metadata_from_memoryview

        rmeta = read_metadata_from_memoryview(memoryview(data_bytes))
    except Exception:
        rmeta = None

    # Prepare result containers
    min_k_hashes: list[list[int]] = []
    histograms: list[list[int]] = []
    min_values: list[int] = []
    null_counts: list[int] = []
    max_values: list[int] = []
    min_values_display: list = []
    max_values_display: list = []

    # iterate schema fields and process each column independently
    schema = pf.schema_arrow
    for col_idx, col_field in enumerate(schema):
        col_name = col_field.name
        try:
            col_table = pf.read(columns=[col_name])
            col = col_table.column(0).combine_chunks()
        except Exception:
            # fallback: try reading the row group column (more granular)
            try:
                tbl = pf.read_row_group(0, columns=[col_name])
                col = tbl.column(0).combine_chunks()
            except Exception:
                # Last resort: read entire file and then take the column
                tbl = pf.read()
                col = tbl.column(col_idx).combine_chunks()

        # compute stats using existing logic encapsulated in helper
        (
            col_min_k,
            col_hist,
            col_min,
            col_max,
            col_min_display,
            col_max_display,
            null_count,
        ) = _compute_stats_for_arrow_column(col, col_field.type, file_path)

        # free the table-level reference if present so memory can be reclaimed
        try:
            del col_table
        except Exception:
            pass
        try:
            del tbl
        except Exception:
            pass

        min_k_hashes.append(col_min_k)
        histograms.append(col_hist)
        min_values.append(col_min)
        max_values.append(col_max)
        min_values_display.append(col_min_display)
        max_values_display.append(col_max_display)
        null_counts.append(null_count)

    # Calculate uncompressed sizes. When the original in-memory table is
    # available (we just wrote it), prefer using it so sizes match the
    # table-based builder exactly. Otherwise materialize the table from
    # bytes and compute sizes the same way.
    import pyarrow as pa
    import pyarrow.parquet as pq

    column_uncompressed: list[int] = []
    uncompressed_size = 0

    # Free references to large objects we no longer need so memory can be reclaimed
    try:
        del buf
    except Exception:
        pass
    try:
        del pf
    except Exception:
        pass
    try:
        del data_bytes
    except Exception:
        pass

    if orig_table is not None:
        # Use the original table buffers so results match the table-based route
        for col in orig_table.columns:
            col_total = 0
            for chunk in col.chunks:
                try:
                    buffs = chunk.buffers()
                except Exception as exc:
                    raise RuntimeError(
                        f"Unable to access chunk buffers to calculate uncompressed size for {file_path}: {exc}"
                    ) from exc
                for buffer in buffs:
                    if buffer is not None:
                        col_total += buffer.size
            column_uncompressed.append(int(col_total))
            uncompressed_size += col_total
    else:
        # Use rugo metadata (if available) to compute per-column uncompressed sizes
        if rmeta:
            rgs = rmeta.get("row_groups", [])
            if rgs:
                ncols = len(rgs[0].get("columns", []))
                for cidx in range(ncols):
                    col_total = 0
                    for rg in rgs:
                        cols = rg.get("columns", [])
                        if cidx < len(cols):
                            col_total += int(cols[cidx].get("total_byte_size", 0) or 0)
                    column_uncompressed.append(int(col_total))
                    uncompressed_size += col_total
                _manifest_metrics["sizes_from_rugo"] += 1
            else:
                column_uncompressed = [0] * len(schema)
                uncompressed_size = 0
                _manifest_metrics["sizes_from_rugo_missing"] += 1
        else:
            # If rugo metadata isn't available, avoid materializing the table;
            # emit zero sizes (safe and memory-light) and track that we lacked
            # metadata for sizes.
            column_uncompressed = [0] * len(schema)
            uncompressed_size = 0
            _manifest_metrics["sizes_from_rugo_unavailable"] += 1
            logger.debug(
                "rugo metadata unavailable for %s; emitting zero column sizes to avoid materializing table",
                file_path,
            )

    entry = ParquetManifestEntry(
        file_path=file_path,
        file_format="parquet",
        record_count=int(meta.num_rows),
        file_size_in_bytes=int(file_size_in_bytes or data_bytes_len),
        uncompressed_size_in_bytes=uncompressed_size,
        column_uncompressed_sizes_in_bytes=column_uncompressed,
        null_counts=null_counts,
        min_k_hashes=min_k_hashes,
        histogram_counts=histograms,
        histogram_bins=HISTOGRAM_BINS,
        min_values=min_values,
        max_values=max_values,
        min_values_display=min_values_display,
        max_values_display=max_values_display,
    )

    logger.debug(
        "build_parquet_manifest_entry_from_bytes %s files=%d dur=%.3fs",
        file_path,
        _manifest_metrics["files_read"],
        time.perf_counter() - t_start,
    )
    return entry


# Backwards-compatible wrapper that keeps the original calling convention
# when a pyarrow Table is already provided (tests and some scripts rely on it).
def build_parquet_manifest_entry(
    table: Any, file_path: str, file_size_in_bytes: int | None = None
) -> ParquetManifestEntry:
    """DEPRECATED: explicit table-based manifest building is removed.

    The implementation previously accepted a PyArrow ``table`` and performed
    the same per-column statistics calculation. That behavior hid a different
    IO/scan path and led to inconsistent performance characteristics.

    Use ``build_parquet_manifest_entry_from_bytes(data_bytes, file_path, file_size_in_bytes, orig_table=None)``
    instead. If you have an in-memory table you can serialize it and call the
    bytes-based builder, or pass ``orig_table`` to preserve exact uncompressed
    size calculations.

    This function now fails fast to avoid silently using the removed path.
    """
    raise RuntimeError(
        "table-based manifest builder removed: use build_parquet_manifest_entry_from_bytes(data_bytes, file_path, file_size_in_bytes, orig_table=table) instead"
    )


def get_manifest_metrics() -> dict:
    """Return a snapshot of manifest instrumentation counters (for tests/benchmarks)."""
    return dict(_manifest_metrics)


def reset_manifest_metrics() -> None:
    """Reset the manifest metrics counters to zero."""
    _manifest_metrics.clear()
