"""Core primitives: AudioModule, AudioPatch, ProvenanceEvent

Small, well-scoped implementations suitable for Phase0 tests.
"""
from __future__ import annotations

import json
import hashlib
import time
import uuid
from dataclasses import dataclass, field, asdict
from typing import List, Dict, Tuple

from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey, Ed25519PublicKey
from cryptography.hazmat.primitives import serialization


def canonical_json(data) -> bytes:
    # deterministic JSON encoding
    return json.dumps(data, separators=(",", ":"), sort_keys=True, ensure_ascii=False).encode("utf-8")


def sha256(b: bytes) -> str:
    return hashlib.sha256(b).hexdigest()


@dataclass
class AudioModule:
    id: str
    type: str
    params: Dict
    version: str = "0.1"
    license: str = ""

    def to_dict(self) -> Dict:
        return asdict(self)

    def cid(self) -> str:
        """Content ID: SHA-256 over canonical JSON of the module."""
        data = self.to_dict()
        return sha256(canonical_json(data))

    @staticmethod
    def make(type: str, params: Dict, version: str = "0.1", license: str = "") -> "AudioModule":
        return AudioModule(id=str(uuid.uuid4()), type=type, params=params, version=version, license=license)


def merkle_tree_root(leaves: List[bytes]) -> str:
    """Simple Merkle tree root (binary) over raw leaf bytes. Returns hex digest."""
    if not leaves:
        return sha256(b"")

    level = [hashlib.sha256(l).digest() for l in leaves]

    while len(level) > 1:
        next_level = []
        for i in range(0, len(level), 2):
            left = level[i]
            right = level[i + 1] if i + 1 < len(level) else left
            next_level.append(hashlib.sha256(left + right).digest())
        level = next_level

    return level[0].hex()


def merkle_proof(leaves: List[bytes], index: int) -> List[Tuple[str, str]]:
    """Return a list of (sibling_hash(hex), direction) where direction is 'L' or 'R'."""
    if index < 0 or index >= len(leaves):
        raise IndexError("leaf index")

    # compute tree levels as bytes
    levels = [ [hashlib.sha256(l).digest() for l in leaves] ]

    while len(levels[-1]) > 1:
        cur = levels[-1]
        nxt = []
        for i in range(0, len(cur), 2):
            left = cur[i]
            right = cur[i+1] if i+1 < len(cur) else left
            nxt.append(hashlib.sha256(left + right).digest())
        levels.append(nxt)

    proof = []
    idx = index
    for level in levels[:-1]:
        sibling_idx = idx ^ 1
        sibling = level[sibling_idx] if sibling_idx < len(level) else level[idx]
        direction = 'L' if sibling_idx < idx else 'R'
        proof.append((sibling.hex(), direction))
        idx = idx // 2

    return proof


def verify_merkle_proof(leaf: bytes, proof: List[Tuple[str, str]], root_hex: str) -> bool:
    cur = hashlib.sha256(leaf).digest()
    for sibling_hex, direction in proof:
        sibling = bytes.fromhex(sibling_hex)
        if direction == 'L':
            cur = hashlib.sha256(sibling + cur).digest()
        else:
            cur = hashlib.sha256(cur + sibling).digest()
    return cur.hex() == root_hex


@dataclass
class AudioPatch:
    nodes: List[Dict] = field(default_factory=list)  # each node: {id, module_cid, params}
    edges: List[Tuple[str, str]] = field(default_factory=list)  # list of (from_node_id, to_node_id)
    metadata: Dict = field(default_factory=dict)

    def to_dict(self) -> Dict:
        return {
            "nodes": self.nodes,
            "edges": self.edges,
            "metadata": self.metadata,
        }

    def merkle_root(self) -> str:
        # leaves are canonical JSON of each node (sorted by node id) and each edge
        node_items = sorted(self.nodes, key=lambda n: n.get("id"))
        node_leaves = [canonical_json(n) for n in node_items]
        edge_items = sorted([{"from": f, "to": t} for f,t in self.edges], key=lambda e: (e['from'], e['to']))
        edge_leaves = [canonical_json(e) for e in edge_items]
        leaves = node_leaves + edge_leaves
        return merkle_tree_root(leaves)

    def inclusion_proof_for_module_cid(self, module_cid: str) -> Tuple[List[Tuple[str,str]], bytes]:
        # find the node(s) that reference this module_cid and return proof for first occurrence
        node_items = sorted(self.nodes, key=lambda n: n.get("id"))
        node_leaves = [canonical_json(n) for n in node_items]
        edge_items = sorted([{"from": f, "to": t} for f,t in self.edges], key=lambda e: (e['from'], e['to']))
        edge_leaves = [canonical_json(e) for e in edge_items]
        leaves = node_leaves + edge_leaves

        for idx, n in enumerate(node_items):
            if n.get("module_cid") == module_cid:
                proof = merkle_proof(leaves, idx)
                return proof, node_leaves[idx]

        raise KeyError("module_cid not referenced in patch")


@dataclass
class ProvenanceEvent:
    patch_id: str
    actor: str
    timestamp: float
    op_hash: str
    signature: bytes = b""

    def to_dict(self) -> Dict:
        return {"patch_id": self.patch_id, "actor": self.actor, "timestamp": self.timestamp, "op_hash": self.op_hash}

    def sign(self, private_key: Ed25519PrivateKey) -> None:
        self.signature = private_key.sign(canonical_json(self.to_dict()))

    def verify(self, public_key: Ed25519PublicKey) -> bool:
        try:
            public_key.verify(self.signature, canonical_json(self.to_dict()))
            return True
        except Exception:
            return False


def generate_ed25519_keypair() -> Tuple[Ed25519PrivateKey, Ed25519PublicKey]:
    sk = Ed25519PrivateKey.generate()
    pk = sk.public_key()
    return sk, pk