"""
File storage adapter for HieraChain Ledger

This module provides a file-based storage implementation for the HieraChain system.
It stores blockchain data in a structured directory layout with separate folders
for different types of data.
"""

import json
import os
import logging
import time
import re
from pathlib import Path
import pyarrow as pa
import pyarrow.parquet as pq
import pyarrow.dataset as ds
import pyarrow.compute as pc

from hierachain.core.block import Block, table_to_list_of_dicts

logger = logging.getLogger(__name__)


def _validate_filename(name: str) -> None:
    """
    Validate filename against strict security rules (CWE-22).
    Allowed: alphanumeric, underscore, hyphen.
    """
    if not re.match(r"^[a-zA-Z0-9_\-]+$", name):
        raise ValueError(
            f"Security: Invalid name '{name}'. "
            "Only alphanumeric, underscore, and hyphen are allowed."
        )


def _get_sorted_block_files(chain_dir: Path) -> list[Path]:
    """Get all block files in a directory sorted by index"""
    return sorted(
        [f for f in chain_dir.glob("block_*.parquet")],
        key=lambda x: int(x.stem.split("_")[1])
    )


def _parse_block_file(block_file: Path) -> dict | None:
    """Helper to read and parse a block Parquet file"""
    try:
        # Read Parquet file
        table = pq.read_table(block_file)

        # Extract metadata
        meta = table.schema.metadata
        if not meta:
            logger.warning("Block file %s missing metadata", block_file)
            return None

        # Decode metadata
        return {
            "index": int(meta.get(b"index", b"0")),
            "events": table,
            "timestamp": float(meta.get(b"timestamp", b"0.0")),
            "previous_hash": meta.get(b"previous_hash", b"").decode("utf-8"),
            "nonce": int(meta.get(b"nonce", b"0")),
            "hash": meta.get(b"hash", b"").decode("utf-8")
        }
    except Exception as e:
        logger.warning("Failed to parse block file %s: %s", block_file, e)
        return None


def _cleanup_chain_blocks(chain_dir: Path, cutoff_time: float):
    """Clean up old block files in a specific chain directory"""
    if not chain_dir.is_dir():
        return

    for block_file in chain_dir.glob("block_*.parquet"):
        try:
            if block_file.stat().st_mtime < cutoff_time:
                block_file.unlink()
                logger.debug("Cleaned up old block file: %s", block_file)
        except OSError as e:
            logger.warning("Failed to delete old block file %s: %s", block_file, e)


def _prepare_block_write_data(block: Block) -> pa.Table:
    """Prepare block data and metadata for storage."""
    # Prepare header metadata
    metadata = {
        b'index': str(block.index).encode(),
        b'timestamp': str(block.timestamp).encode(),
        b'previous_hash': str(block.previous_hash).encode(),
        b'nonce': str(block.nonce).encode(),
        b'hash': str(block.hash).encode(),
        b'stored_at': str(time.time()).encode()
    }

    # Merge with existing schema metadata
    table = block.events
    existing_meta = table.schema.metadata or {}
    return table.replace_schema_metadata({**existing_meta, **metadata})


def _write_parquet_optimized(table: pa.Table, file_path: Path):
    """Write Arrow table to Parquet with optimized settings."""
    pq.write_table(
        table,
        file_path,
        compression="zstd",
        compression_level=3,
        use_dictionary=True,
        write_statistics=True
    )


def _build_event_index_table(events_list: list, block_data: dict) -> pa.Table | None:
    """Extract event data and build an index table."""
    indices = []
    for event in events_list:
        entity_id = event.get("entity_id")
        if not entity_id:
            continue

        indices.append({
            "entity_id": entity_id,
            "block_index": block_data["index"],
            "event_type": event.get("event", event.get("event_type", "unknown")),
            "timestamp": event.get("timestamp", block_data["timestamp"]),
            "block_hash": block_data["hash"]
        })

    if not indices:
        return None

    schema = pa.schema([
        ("entity_id", pa.string()),
        ("block_index", pa.int64()),
        ("event_type", pa.string()),
        ("timestamp", pa.float64()),
        ("block_hash", pa.string())
    ])

    pydict = {name: [row.get(name) for row in indices] for name in schema.names}
    return pa.table(pydict, schema=schema)


def _count_blocks(chain_dir: Path) -> int:
    """Count block files in directory."""
    return len(list(chain_dir.glob("block_*.parquet")))


def _get_dataset_stats(events_dir: Path) -> tuple[int, int]:
    """Get row count and unique entity count for a dataset."""
    if not events_dir.exists() or not any(events_dir.iterdir()):
        return 0, 0

    try:
        dataset = ds.dataset(events_dir, format="parquet")
        total_events = dataset.count_rows()
        unique_table = dataset.scanner(columns=["entity_id"]).to_table()
        unique_count = len(unique_table.column("entity_id").unique())
        return total_events, unique_count
    except Exception as e:
        logger.warning("Failed to load dataset stats: %s", e)
        return 0, 0


def _calculate_storage_stats(storage_path: Path) -> tuple[int, int]:
    """Calculate total size and file count for a directory."""
    total_size = 0
    file_count = 0
    for root, _, files in os.walk(storage_path):
        for file in files:
            file_path = Path(root) / file
            total_size += file_path.stat().st_size
            file_count += 1
    return total_size, file_count


def _filter_block_events(
    events_table: pa.Table,
    entity_id: str,
    record_timestamp: float
) -> list[dict]:
    """Filter Arrow table for specific entity and timestamp."""
    expr = (
        (pc.field("entity_id") == entity_id) &
        (pc.field("timestamp") == record_timestamp)
    )
    subset = events_table.filter(expr)
    return table_to_list_of_dicts(subset)


def _map_event_record(
    full_event: dict,
    search_chain: str,
    block_index: int
) -> dict:
    """Map raw event record to standard output format."""
    return {
        "chain_name": search_chain,
        "block_index": block_index,
        "event_type": full_event.get("event"),
        "timestamp": full_event.get("timestamp"),
        "details": full_event.get("details", {})
    }


def _get_chain_events_dataset(events_dir: Path, entity_id: str) -> list[dict]:
    """Load and scan events dataset for an entity."""
    dataset = ds.dataset(events_dir, format="parquet")
    scanner = dataset.scanner(filter=pc.field("entity_id") == entity_id)
    filtered_table = scanner.to_table()
    sorted_table = filtered_table.sort_by([("timestamp", "ascending")])
    return sorted_table.to_pylist()


def _load_block_files_batch(block_files: list[Path]) -> list[dict]:
    """Parse a batch of block files and return valid data."""
    blocks = []
    for block_file in block_files:
        block_data = _parse_block_file(block_file)
        if block_data:
            blocks.append(block_data)
    return blocks


def _fetch_events_from_chain(
    events_dir: Path,
    storage: "FileStorageAdapter",
    search_chain: str,
    entity_id: str
) -> list[dict]:
    """Helper to fetch events from a chain for an entity."""
    chain_events = []

    try:
        index_records = _get_chain_events_dataset(events_dir, entity_id)
        for record in index_records:
            block = storage.get_block(search_chain, record["block_index"])
            if not block:
                continue

            rows = _filter_block_events(block["events"], entity_id, record["timestamp"])
            for row in rows:
                chain_events.append(
                    _map_event_record(row, search_chain, record["block_index"])
                )
    except Exception as e:
        logger.warning("Failed to load events for chain %s: %s", search_chain, e)

    return chain_events


def _fetch_optimized_events(
    events_dir: Path,
    search_chain: str,
    entity_id: str,
    projection: list[str] | None = None
) -> list[dict]:
    """Helper for optimized event fetching using Arrow scanner."""
    chain_events = []
    try:
        dataset = ds.dataset(events_dir, format="parquet")
        scanner = dataset.scanner(
            filter=pc.field("entity_id") == entity_id,
            columns=projection
        )
        filtered_table = scanner.to_table()
        sorted_table = filtered_table.sort_by([("timestamp", "ascending")])

        for record in sorted_table.to_pylist():
            record["chain_name"] = search_chain
            chain_events.append(record)
    except Exception as e:
        logger.error("Failed to process chain %s: %s", search_chain, e)

    return chain_events


def _execute_event_search(
    chains: list[str],
    entity_id: str,
    fetch_fn
) -> list[dict]:
    """Search and sort events across multiple chains."""
    events = []
    for chain in chains:
        events.extend(fetch_fn(chain, entity_id))
    events.sort(key=lambda x: x.get("timestamp", 0))
    return events


def _run_data_cleanup(blocks_path: Path, days_to_keep: int):
    """Cleanup old data files across all chains."""
    cutoff_time = time.time() - (days_to_keep * 24 * 60 * 60)
    for chain_dir in blocks_path.iterdir():
        if chain_dir.is_dir():
            _cleanup_chain_blocks(chain_dir, cutoff_time)


class FileStorageAdapter:
    """File-based storage adapter for blockchain data"""

    def __init__(self, storage_path: str = "blockchain_data"):
        """
        Initialize file storage adapter

        Args:
            storage_path: Base directory for storing blockchain data
        """
        self.storage_path = Path(storage_path)
        self.chains_path = self.storage_path / "chains"
        self.blocks_path = self.storage_path / "blocks"
        self.events_path = self.storage_path / "events"
        self.proofs_path = self.storage_path / "proofs"

        self._create_directories()
        logger.info("File storage initialized at: %s", self.storage_path)

    def _create_directories(self):
        """Create necessary directories for storage"""
        directories = [
            self.storage_path,
            self.chains_path,
            self.blocks_path,
            self.events_path,
            self.proofs_path
        ]

        for directory in directories:
            directory.mkdir(parents=True, exist_ok=True)

    def _get_chain_file(self, chain_name: str) -> Path:
        """Get file path for chain metadata"""
        _validate_filename(chain_name)
        return self.chains_path / f"{chain_name}.json"

    def _get_block_file(self, chain_name: str, block_index: int) -> Path:
        """Get file path for a specific block (Parquet)"""
        _validate_filename(chain_name)
        chain_dir = self.blocks_path / chain_name
        chain_dir.mkdir(exist_ok=True)
        return chain_dir / f"block_{block_index:06d}.parquet"

    def _get_events_dir(self, chain_name: str) -> Path:
        """Get directory for chain events dataset"""
        _validate_filename(chain_name)
        path = self.events_path / chain_name
        path.mkdir(parents=True, exist_ok=True)
        return path

    def store_chain_metadata(
        self,
        chain_name: str,
        chain_type: str,
        parent_chain: str = None,
        metadata: dict = None
    ):
        """Store chain metadata"""
        try:
            chain_data = {
                "name": chain_name,
                "type": chain_type,
                "parent_chain": parent_chain,
                "metadata": metadata or {},
                "created_at": time.time(),
                "updated_at": time.time()
            }

            chain_file = self._get_chain_file(chain_name)
            with open(chain_file, "w", encoding="utf-8") as f:
                json.dump(chain_data, f, indent=2)

            logger.debug("Stored chain metadata: %s", chain_name)

        except Exception as e:
            logger.error("Failed to store chain metadata %s: %s", chain_name, e)
            raise

    def store_block(self, chain_name: str, block_data: dict | Block):
        """
        Store block data using Parquet for high performance.

        Args:
            chain_name: Name of the chain
            block_data: Block object or Dictionary
        """
        try:
            # Convert to Block object if it's a dict
            block = (
                Block.from_dict(block_data)
                if isinstance(block_data, dict) else block_data
            )

            # Prepare data and write
            table_with_meta = _prepare_block_write_data(block)
            block_file = self._get_block_file(chain_name, block.index)
            _write_parquet_optimized(table_with_meta, block_file)

            self._update_events_index(chain_name, block.to_dict())
            logger.debug("Stored block %s for chain %s", block.index, chain_name)

        except Exception as e:
            logger.error("Failed to store block: %s", e)
            raise

    def _update_events_index(self, chain_name: str, block_data: dict):
        """Update events index using Arrow Dataset."""
        try:
            events_list = block_data.get("events", [])
            if not events_list:
                return

            table = _build_event_index_table(events_list, block_data)
            if table is None:
                return

            events_dir = self._get_events_dir(chain_name)
            file_path = events_dir / f"events_{block_data['index']:09d}.parquet"
            _write_parquet_optimized(table, file_path)

        except Exception as e:
            logger.error("Failed to update events index: %s", e)

    def get_chain_metadata(self, chain_name: str) -> dict | None:
        """Get chain metadata"""
        try:
            chain_file = self._get_chain_file(chain_name)
            if not chain_file.exists():
                return None

            with open(chain_file, 'r', encoding='utf-8') as f:
                return json.load(f)
    
        except Exception as e:
            logger.error("Failed to get chain metadata %s: %s", chain_name, e)
            return None

    def get_block(self, chain_name: str, block_index: int) -> dict | None:
        """
        Get a specific block.
        Returns a Dictionary suitable for Block.from_dict(),
        with 'events' as a pyarrow.Table.
        """
        try:
            block_file = self._get_block_file(chain_name, block_index)
            if not block_file.exists():
                return None

            return _parse_block_file(block_file)

        except Exception as e:
            logger.error(
                "Failed to get block %s for chain %s: %s",
                block_index, chain_name, e
            )
            return None

    def get_chain_blocks(
        self, chain_name: str, limit: int | None = None, offset: int = 0
    ) -> list[dict]:
        """Get blocks for a specific chain"""
        try:
            _validate_filename(chain_name)
            chain_dir = self.blocks_path / chain_name
            if not chain_dir.exists():
                return []

            block_files = _get_sorted_block_files(chain_dir)
            start = offset
            end = (offset + limit) if limit is not None else None
            return _load_block_files_batch(block_files[start:end])

        except ValueError:
            raise
        except Exception as e:
            logger.error("Failed to get blocks for chain %s: %s", chain_name, e)
            return []

    def _get_search_chains(self, chain_name: str | None = None) -> list[str]:
        """Determine which chains to search for events"""
        if chain_name:
            return [chain_name]
        return self.list_chains()

    def _get_events_from_chain(self, search_chain: str, entity_id: str) -> list[dict]:
        """Get events for an entity from a specific chain's events dataset"""
        events_dir = self._get_events_dir(search_chain)
        return _fetch_events_from_chain(events_dir, self, search_chain, entity_id)

    def get_entity_events(
        self, entity_id: str, chain_name: str | None = None
    ) -> list[dict]:
        """Get all events for a specific entity using Arrow Dataset."""
        try:
            chains = self._get_search_chains(chain_name)
            return _execute_event_search(chains, entity_id, self._get_events_from_chain)
        except ValueError:
            raise
        except Exception as e:
            logger.error("Failed to get events for entity %s: %s", entity_id, e)
            return []

    def get_chain_stats(self, chain_name: str) -> dict:
        """Get statistics for a specific chain."""
        try:
            _validate_filename(chain_name)
            chain_dir = self.blocks_path / chain_name
            events_dir = self._get_events_dir(chain_name)

            if not chain_dir.exists():
                return {
                    "chain_name": chain_name,
                    "total_blocks": 0,
                    "total_events": 0,
                    "unique_entities": 0
                }

            total_blocks = _count_blocks(chain_dir)
            total_events, unique_entities = _get_dataset_stats(events_dir)

            return {
                "chain_name": chain_name,
                "total_blocks": total_blocks,
                "total_events": total_events,
                "unique_entities": unique_entities
            }

        except Exception as e:
            if isinstance(e, ValueError): raise
            logger.error("Failed to get stats for chain %s: %s", chain_name, e)
            return {
                "chain_name": chain_name,
                "total_blocks": 0,
                "total_events": 0,
                "unique_entities": 0
            }

    def list_chains(self) -> list[str]:
        """list all stored chains"""
        try:
            chain_files = self.chains_path.glob("*.json")
            return [f.stem for f in chain_files]
        except Exception as e:
            logger.error("Failed to list chains: %s", e)
            return []

    def cleanup_old_data(self, days_to_keep: int = 30):
        """Clean up old data files"""
        try:
            _run_data_cleanup(self.blocks_path, days_to_keep)
            logger.info("Cleaned up data older than %s days", days_to_keep)
        except Exception as e:
            logger.error("Failed to cleanup old data: %s", e)

    def get_storage_info(self) -> dict:
        """Get storage information"""
        try:
            total_size, file_count = _calculate_storage_stats(self.storage_path)

            return {
                "storage_path": str(self.storage_path),
                "total_size_bytes": total_size,
                "total_size_mb": round(total_size / (1024 * 1024), 2),
                "file_count": file_count,
                "chains_count": len(self.list_chains())
            }

        except Exception as e:
            logger.error("Failed to get storage info: %s", e)
            return {}

    def _process_optimized_chain_events(
        self,
        search_chain: str,
        entity_id: str,
        projection: list[str] | None = None
    ) -> list[dict]:
        """Process events for a chain in an optimized way"""
        events_dir = self._get_events_dir(search_chain)
        return _fetch_optimized_events(events_dir, search_chain, entity_id, projection)

    def get_entity_events_optimized(
        self,
        entity_id: str,
        chain_name: str | None = None,
        columns: list[str] | None = None
    ) -> list[dict]:
        """Get events with column pruning for better performance."""
        try:
            chains = self._get_search_chains(chain_name)
            projection = columns if columns else None
            
            def fetch_fn(c, eid):
                return self._process_optimized_chain_events(c, eid, projection)
                
            return _execute_event_search(chains, entity_id, fetch_fn)
        except Exception as e:
            logger.error("Failed to get optimized events for %s: %s", entity_id, e)
            return []


class BatchBlockWriter:
    """
    Optimized batch writer for multiple blocks.

    Uses buffering for better I/O performance when writing
    multiple blocks in sequence.

    Usage:
        with BatchBlockWriter(storage, 'my_chain', batch_size=50) as writer:
            for block in blocks:
                writer.add(block)
    """

    def __init__(
        self,
        storage: FileStorageAdapter,
        chain_name: str,
        batch_size: int = 100
    ):
        """
        Initialize batch writer.

        Args:
            storage: FileStorageAdapter instance
            chain_name: Name of chain to write to
            batch_size: Number of blocks to buffer before flush
        """
        self.storage = storage
        self.chain_name = chain_name
        self.batch_size = batch_size
        self._buffer: list[Block] = []
        self._stats = {
            "blocks_written": 0,
            "events_written": 0,
            "flush_count": 0,
            "total_time_ms": 0.0
        }

    def add(self, block: Block) -> None:
        """
        Add block to buffer.

        Args:
            block: Block to add
        """
        self._buffer.append(block)
        if len(self._buffer) >= self.batch_size:
            self.flush()

    def flush(self) -> None:
        """Write all buffered blocks to storage."""
        if not self._buffer:
            return

        start_ts = time.time()
        for block in self._buffer:
            self.storage.store_block(self.chain_name, block)
            self._stats["blocks_written"] += 1
            self._stats["events_written"] += len(block.events)

        elapsed_ms = (time.time() - start_ts) * 1000
        self._stats["flush_count"] += 1
        self._stats["total_time_ms"] += elapsed_ms
        logger.debug("Flushed %s blocks in %.2fms", len(self._buffer), elapsed_ms)
        self._buffer.clear()

    def get_stats(self) -> dict:
        """
        Get write statistics.

        Returns:
            Dictionary with blocks_written, events_written,
            flush_count, total_time_ms, avg_time_per_block_ms
        """
        stats = self._stats.copy()
        if stats["blocks_written"] > 0:
            stats["avg_time_per_block_ms"] = (
                stats["total_time_ms"] / stats["blocks_written"]
            )
        else:
            stats["avg_time_per_block_ms"] = 0.0
        return stats

    def __enter__(self) -> 'BatchBlockWriter':
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        self.flush()
