"""Merkle tree integrity verification for resilient storage.

Provides cryptographic integrity verification through Merkle trees,
allowing efficient proof of shard inclusion and data integrity without
needing to download all shards.
"""

from __future__ import annotations

import hashlib
from collections.abc import Sequence
from dataclasses import dataclass, field
from typing import Any

from .models import (
    IntegrityReport,
    ShardSet,
    StorageShard,
)


def compute_hash(data: bytes, algorithm: str = "sha256") -> str:
    """Compute cryptographic hash of data.

    Args:
        data: Bytes to hash
        algorithm: Hash algorithm (default: sha256)

    Returns:
        Hex-encoded hash string
    """
    if algorithm == "sha256":
        return hashlib.sha256(data).hexdigest()
    elif algorithm == "sha3_256":
        return hashlib.sha3_256(data).hexdigest()
    elif algorithm == "blake2b":
        return hashlib.blake2b(data, digest_size=32).hexdigest()
    else:
        raise ValueError(f"Unsupported hash algorithm: {algorithm}")


def _hash_pair(left: str, right: str, algorithm: str = "sha256") -> str:
    """Hash a pair of hex strings together."""
    combined = bytes.fromhex(left) + bytes.fromhex(right)
    return compute_hash(combined, algorithm)


@dataclass
class MerkleProof:
    """Proof of inclusion in a Merkle tree.

    Contains the path from a leaf to the root, allowing verification
    that a specific piece of data is part of the tree without needing
    the entire dataset.
    """

    leaf_hash: str  # Hash of the leaf data
    leaf_index: int  # Position in the tree (0-indexed)
    proof_hashes: list[str] = field(default_factory=list)  # Sibling hashes on path to root
    proof_directions: list[bool] = field(default_factory=list)  # True = sibling is on right
    root_hash: str = ""  # Expected root hash

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary."""
        return {
            "leaf_hash": self.leaf_hash,
            "leaf_index": self.leaf_index,
            "proof_hashes": self.proof_hashes,
            "proof_directions": self.proof_directions,
            "root_hash": self.root_hash,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> MerkleProof:
        """Create from dictionary."""
        return cls(
            leaf_hash=data["leaf_hash"],
            leaf_index=data["leaf_index"],
            proof_hashes=data.get("proof_hashes", []),
            proof_directions=data.get("proof_directions", []),
            root_hash=data.get("root_hash", ""),
        )


def verify_proof(proof: MerkleProof, algorithm: str = "sha256") -> bool:
    """Verify a Merkle proof.

    Args:
        proof: The MerkleProof to verify
        algorithm: Hash algorithm used in the tree

    Returns:
        True if the proof is valid (leaf hash leads to root hash)
    """
    if not proof.proof_hashes:
        # Single node tree - leaf is root
        return proof.leaf_hash == proof.root_hash

    current = proof.leaf_hash

    for sibling_hash, sibling_right in zip(proof.proof_hashes, proof.proof_directions, strict=True):
        if sibling_right:
            # Sibling is on the right
            current = _hash_pair(current, sibling_hash, algorithm)
        else:
            # Sibling is on the left
            current = _hash_pair(sibling_hash, current, algorithm)

    return current == proof.root_hash


@dataclass
class MerkleTree:
    """Merkle tree for integrity verification.

    Builds a binary hash tree from leaf data, providing:
    - O(1) root hash for integrity verification
    - O(log n) proof of inclusion for any leaf
    - Efficient detection of corrupted nodes

    Example:
        tree = MerkleTree.from_data([shard1, shard2, shard3])
        root = tree.root_hash
        proof = tree.get_proof(0)
        assert verify_proof(proof)
    """

    leaves: list[str] = field(default_factory=list)  # Leaf hashes
    levels: list[list[str]] = field(default_factory=list)  # All tree levels
    algorithm: str = "sha256"

    @property
    def root_hash(self) -> str:
        """Get the root hash of the tree."""
        if not self.levels:
            return ""
        return self.levels[-1][0] if self.levels[-1] else ""

    @property
    def height(self) -> int:
        """Height of the tree (number of levels)."""
        return len(self.levels)

    @property
    def leaf_count(self) -> int:
        """Number of leaves in the tree."""
        return len(self.leaves)

    @classmethod
    def from_data(cls, data_items: Sequence[bytes], algorithm: str = "sha256") -> MerkleTree:
        """Build a Merkle tree from raw data items.

        Args:
            data_items: List of bytes to include as leaves
            algorithm: Hash algorithm to use

        Returns:
            Constructed MerkleTree
        """
        if not data_items:
            return cls(leaves=[], levels=[[]], algorithm=algorithm)

        # Hash all leaves
        leaves = [compute_hash(data, algorithm) for data in data_items]

        return cls.from_hashes(leaves, algorithm)

    @classmethod
    def from_hashes(cls, leaf_hashes: list[str], algorithm: str = "sha256") -> MerkleTree:
        """Build a Merkle tree from pre-computed leaf hashes.

        Args:
            leaf_hashes: List of hex-encoded leaf hashes
            algorithm: Hash algorithm to use for internal nodes

        Returns:
            Constructed MerkleTree
        """
        if not leaf_hashes:
            return cls(leaves=[], levels=[[]], algorithm=algorithm)

        leaves = list(leaf_hashes)
        levels = [leaves]

        # Build tree bottom-up
        current_level = leaves
        while len(current_level) > 1:
            next_level = []

            # Process pairs
            for i in range(0, len(current_level), 2):
                left = current_level[i]
                # If odd number of nodes, duplicate the last one
                right = current_level[i + 1] if i + 1 < len(current_level) else left
                parent = _hash_pair(left, right, algorithm)
                next_level.append(parent)

            levels.append(next_level)
            current_level = next_level

        return cls(leaves=leaves, levels=levels, algorithm=algorithm)

    @classmethod
    def from_shards(cls, shards: Sequence[StorageShard], algorithm: str = "sha256") -> MerkleTree:
        """Build a Merkle tree from storage shards.

        Args:
            shards: List of StorageShard objects
            algorithm: Hash algorithm to use

        Returns:
            Constructed MerkleTree
        """
        # Use existing checksums if available, otherwise compute
        leaf_hashes = []
        for shard in shards:
            if shard.metadata.checksum:
                leaf_hashes.append(shard.metadata.checksum)
            else:
                leaf_hashes.append(compute_hash(shard.data, algorithm))

        return cls.from_hashes(leaf_hashes, algorithm)

    def get_proof(self, leaf_index: int) -> MerkleProof:
        """Generate a proof of inclusion for a leaf.

        Args:
            leaf_index: Index of the leaf (0-indexed)

        Returns:
            MerkleProof that can be used to verify inclusion

        Raises:
            IndexError: If leaf_index is out of range
        """
        if leaf_index < 0 or leaf_index >= len(self.leaves):
            raise IndexError(f"Leaf index {leaf_index} out of range [0, {len(self.leaves)})")

        proof_hashes = []
        proof_directions = []

        current_index = leaf_index

        for level in self.levels[:-1]:  # Skip root level
            # Determine sibling index
            if current_index % 2 == 0:
                # Current is left, sibling is right
                sibling_index = current_index + 1
                sibling_right = True
            else:
                # Current is right, sibling is left
                sibling_index = current_index - 1
                sibling_right = False

            # Get sibling hash (handle odd-length levels)
            if sibling_index < len(level):
                proof_hashes.append(level[sibling_index])
            else:
                # No sibling (odd level), use current node as its own sibling
                proof_hashes.append(level[current_index])

            proof_directions.append(sibling_right)

            # Move to parent index
            current_index //= 2

        return MerkleProof(
            leaf_hash=self.leaves[leaf_index],
            leaf_index=leaf_index,
            proof_hashes=proof_hashes,
            proof_directions=proof_directions,
            root_hash=self.root_hash,
        )

    def verify_leaf(self, leaf_index: int, data: bytes) -> bool:
        """Verify that data at leaf_index is correct.

        Args:
            leaf_index: Index of the leaf
            data: Data to verify

        Returns:
            True if data matches the leaf hash
        """
        if leaf_index < 0 or leaf_index >= len(self.leaves):
            return False

        data_hash = compute_hash(data, self.algorithm)
        return data_hash == self.leaves[leaf_index]

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary."""
        return {
            "leaves": self.leaves,
            "root_hash": self.root_hash,
            "height": self.height,
            "algorithm": self.algorithm,
        }


@dataclass
class IntegrityVerifier:
    """Verifies integrity of shard sets using Merkle trees.

    Provides comprehensive integrity checking including:
    - Individual shard checksum verification
    - Merkle tree consistency
    - Recoverability assessment
    """

    algorithm: str = "sha256"

    def verify_shard_set(self, shard_set: ShardSet) -> IntegrityReport:
        """Perform full integrity verification on a shard set.

        Args:
            shard_set: ShardSet to verify

        Returns:
            IntegrityReport with detailed results
        """
        total_shards = shard_set.total_shards_n
        available_shards = shard_set.available_shards

        # Check individual shard checksums
        valid_shards = []
        corrupted_indices = []
        missing_indices = list(shard_set.missing_indices)

        for shard in available_shards:
            if shard.is_valid:
                valid_shards.append(shard)
            else:
                corrupted_indices.append(shard.index)

        valid_count = len(valid_shards)

        # Check Merkle tree if root is available
        merkle_valid = True
        if shard_set.merkle_root and valid_shards:
            # Rebuild tree from available shards and verify root
            tree = MerkleTree.from_shards(valid_shards, self.algorithm)
            # For partial trees, we can't fully verify - just check structure
            merkle_valid = len(tree.leaves) == len(valid_shards)

        # Check if original data checksum would match (requires recovery)
        checksum_valid = True  # Assumed true unless we recover and check

        # Determine if data is recoverable
        can_recover = valid_count >= shard_set.data_shards_k

        # Overall validity
        is_valid = len(corrupted_indices) == 0 and len(missing_indices) == 0 and merkle_valid

        # Build details message
        details_parts = []
        if corrupted_indices:
            details_parts.append(f"Corrupted shards: {corrupted_indices}")
        if missing_indices:
            details_parts.append(f"Missing shards: {missing_indices}")
        if not can_recover:
            details_parts.append(f"Cannot recover: need {shard_set.data_shards_k}, have {valid_count}")
        if is_valid:
            details_parts.append("All shards valid and complete")

        return IntegrityReport(
            is_valid=is_valid,
            total_shards=total_shards,
            valid_shards=valid_count,
            corrupted_shards=corrupted_indices,
            missing_shards=missing_indices,
            merkle_valid=merkle_valid,
            checksum_valid=checksum_valid,
            can_recover=can_recover,
            details="; ".join(details_parts),
        )

    def verify_shard(self, shard: StorageShard) -> bool:
        """Verify a single shard's integrity.

        Args:
            shard: StorageShard to verify

        Returns:
            True if shard checksum is valid
        """
        return shard.is_valid

    def generate_merkle_root(self, shard_set: ShardSet) -> str:
        """Generate Merkle root for a shard set.

        Args:
            shard_set: ShardSet to generate root for

        Returns:
            Hex-encoded Merkle root hash
        """
        if not shard_set.shards:
            return ""

        tree = MerkleTree.from_shards(shard_set.available_shards, self.algorithm)
        return tree.root_hash

    def generate_proof(self, shard_set: ShardSet, shard_index: int) -> MerkleProof:
        """Generate a Merkle proof for a specific shard.

        Args:
            shard_set: ShardSet containing the shard
            shard_index: Index of the shard

        Returns:
            MerkleProof for the specified shard
        """
        tree = MerkleTree.from_shards(shard_set.available_shards, self.algorithm)

        # Find the position in available shards
        position = 0
        for shard in shard_set.available_shards:
            if shard.index == shard_index:
                break
            position += 1
        else:
            raise ValueError(f"Shard {shard_index} not found in shard set")

        return tree.get_proof(position)

    def challenge_response_verify(self, shard: StorageShard, expected_checksum: str) -> bool:
        """Challenge-response verification for a shard.

        Used for remote verification where we send a challenge
        and expect a specific response.

        Args:
            shard: Shard to verify
            expected_checksum: Expected SHA-256 checksum

        Returns:
            True if shard data matches expected checksum
        """
        actual = compute_hash(shard.data, self.algorithm)
        return actual == expected_checksum
