"""
HASHAB implementation for iPod Nano 6G and 7G.

Uses a WebAssembly module (calcHashAB.wasm) from dstaley/hashab — a clean-room
reimplementation of Apple's white-box AES signing algorithm.

The output is a 57-byte signature used in two places:
1. iTunesDB binary mhbd header at offset 0xAB
2. Locations.itdb.cbk header (first 57 bytes)

Source: https://github.com/dstaley/hashab (The Unlicense)
"""

import hashlib
import logging
from pathlib import Path

logger = logging.getLogger(__name__)

HASHAB_SIZE = 57

# Path to the WASM module (shipped alongside this file)
_WASM_DIR = Path(__file__).parent / "wasm"
_WASM_PATH = _WASM_DIR / "calcHashAB.wasm"

# Lazy-loaded WASM engine (expensive to create — reuse across calls)
_wasm_instance = None
_wasm_store = None


def _get_wasm_instance():
    """Load the WASM module and return (store, instance)."""
    global _wasm_instance, _wasm_store

    if _wasm_instance is not None:
        return _wasm_store, _wasm_instance

    try:
        import wasmtime
    except ImportError:
        raise ImportError(
            "wasmtime is required for HASHAB (iPod Nano 6G/7G). "
            "Install with: pip install wasmtime"
        )

    if not _WASM_PATH.exists():
        raise FileNotFoundError(
            f"WASM module not found at {_WASM_PATH}. "
            "Download calcHashAB.wasm from "
            "https://github.com/dstaley/hashab/releases/tag/2025-01-04"
        )

    engine = wasmtime.Engine()
    store = wasmtime.Store(engine)
    module = wasmtime.Module.from_file(engine, str(_WASM_PATH))
    instance = wasmtime.Instance(store, module, [])

    _wasm_store = store
    _wasm_instance = instance

    logger.debug("HASHAB WASM module loaded from %s", _WASM_PATH)
    return store, instance


def compute_hashab(sha1_digest: bytes, uuid: bytes) -> bytes:
    """
    Compute 57-byte HASHAB signature using the WASM module.

    Args:
        sha1_digest: 20-byte SHA1 hash
        uuid: 8-byte FireWire GUID from SysInfo

    Returns:
        57-byte signature
    """
    if len(sha1_digest) != 20:
        raise ValueError(f"SHA1 must be 20 bytes, got {len(sha1_digest)}")
    if len(uuid) < 8:
        raise ValueError(f"UUID must be at least 8 bytes, got {len(uuid)}")

    store, instance = _get_wasm_instance()
    exports = instance.exports(store)
    memory = exports["memory"]
    get_input_sha1 = exports["getInputSha1"]
    get_input_uuid = exports["getInputUuid"]
    get_output = exports["getOutput"]
    calculate_hash = exports["calculateHash"]

    sha1_ptr = get_input_sha1(store)
    uuid_ptr = get_input_uuid(store)
    output_ptr = get_output(store)

    mem_data = memory.data_ptr(store)

    for i in range(20):
        mem_data[sha1_ptr + i] = sha1_digest[i]

    for i in range(8):
        mem_data[uuid_ptr + i] = uuid[i]

    calculate_hash(store)

    result = bytes(mem_data[output_ptr + i] for i in range(HASHAB_SIZE))
    return result


def parse_firewire_guid(sysinfo_path: str) -> bytes:
    """Parse FirewireGuid from SysInfo file.
    
    SysInfo contains lines like:
        FirewireGuid: 0x000A270024A41BB7
    
    Returns:
        8-byte FireWire GUID
    """
    with open(sysinfo_path, 'r') as f:
        for line in f:
            if line.startswith('FirewireGuid:'):
                hex_str = line.split(':')[1].strip()
                if hex_str.startswith('0x'):
                    hex_str = hex_str[2:]
                return bytes.fromhex(hex_str)
    raise ValueError(f"No FirewireGuid found in {sysinfo_path}")
