"""Core anchoring functions for W&B artifacts."""

import hashlib
import os
import sys
from typing import Optional

from contextlib import contextmanager


def _hash_file(path: str) -> str:
    """Compute SHA-256 hash of a file. Bytes never leave the machine."""
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(8192), b""):
            h.update(chunk)
    return h.hexdigest()


def _get_client(api_key: Optional[str] = None):
    """Lazy-import UmariseCore to avoid import-time dependency."""
    from umarise import UmariseCore
    key = api_key or os.environ.get("UMARISE_API_KEY")
    if not key:
        raise RuntimeError(
            "UMARISE_API_KEY not set. Export it or pass api_key= to enable()."
        )
    return UmariseCore(api_key=key)


def anchor_file(
    file_path: str,
    run=None,
    api_key: Optional[str] = None,
) -> Optional[str]:
    """
    Anchor a single file to Bitcoin via Umarise.

    Returns:
        origin_id if successful, None otherwise.
    """
    try:
        file_hash = _hash_file(file_path)
        client = _get_client(api_key)
        result = client.attest(f"sha256:{file_hash}")

        origin_id = getattr(result, "origin_id", None) or "unknown"
        proof_status = getattr(result, "proof_status", None) or "pending"

        # Tag the W&B run if available
        if run is not None:
            try:
                run.summary["umarise_anchored"] = True
                run.summary["umarise_origin_id"] = origin_id
                run.summary["umarise_hash"] = file_hash
                run.summary["umarise_proof_status"] = proof_status
            except Exception:
                pass

        print(
            f"[umarise] ✓ Anchored {os.path.basename(file_path)} → {origin_id}",
            file=sys.stderr,
        )
        return origin_id

    except Exception as e:
        print(
            f"[umarise] ⚠ Anchoring failed for {file_path}: {e}",
            file=sys.stderr,
        )
        return None


def anchor_artifact(artifact, run=None, api_key: Optional[str] = None) -> list:
    """
    Anchor all files in a W&B Artifact to Bitcoin.

    Returns:
        List of (file_path, origin_id) tuples.
    """
    results = []

    try:
        manifest = artifact.manifest
        for entry in manifest.entries.values():
            local_path = getattr(entry, "local_path", None)
            if local_path and os.path.isfile(local_path):
                origin_id = anchor_file(local_path, run=run, api_key=api_key)
                results.append((local_path, origin_id))
    except Exception:
        pass

    return results


def anchor_logged_artifact(
    artifact_name: str,
    version: str = "latest",
    api_key: Optional[str] = None,
) -> list:
    """Download and anchor a previously logged W&B artifact."""
    import wandb

    if not wandb.run:
        raise RuntimeError("No active W&B run. Call wandb.init() first.")

    ref = f"{artifact_name}:{version}" if ":" not in artifact_name else artifact_name
    artifact = wandb.run.use_artifact(ref)
    artifact_dir = artifact.download()

    results = []
    for root, _, files in os.walk(artifact_dir):
        for fname in files:
            fpath = os.path.join(root, fname)
            origin_id = anchor_file(fpath, run=wandb.run, api_key=api_key)
            results.append((fpath, origin_id))

    return results
