"""Auto-anchor callback for W&B."""

import os
import sys
import hashlib
import threading
from typing import Optional


# Global state for enable/disable
_original_log_artifact = None
_enabled = False


def enable(api_key: Optional[str] = None):
    """
    Enable automatic anchoring for all wandb.log_artifact() calls.

    Call once at the start of your script:
        import umarise_wandb
        umarise_wandb.enable()

    Every subsequent artifact logged via wandb will be anchored to Bitcoin.
    """
    global _original_log_artifact, _enabled

    if _enabled:
        return

    import wandb
    from umarise_wandb.anchor import _get_client

    client = _get_client(api_key)
    _original_log_artifact = wandb.Run.log_artifact

    def _patched_log_artifact(self, artifact, *args, **kwargs):
        result = _original_log_artifact(self, artifact, *args, **kwargs)

        # Anchor in background after logging
        try:
            if hasattr(artifact, "manifest") and artifact.manifest:
                for entry in artifact.manifest.entries.values():
                    local_path = getattr(entry, "local_path", None)
                    if local_path and os.path.exists(local_path):
                        h = hashlib.sha256()
                        with open(local_path, "rb") as f:
                            for chunk in iter(lambda: f.read(8192), b""):
                                h.update(chunk)
                        file_hash = h.hexdigest()
                        anchor_result = client.attest(f"sha256:{file_hash}")

                        origin_id = getattr(anchor_result, "origin_id", None) or "unknown"
                        if self:
                            try:
                                self.summary["umarise_anchored"] = True
                                self.summary["umarise_origin_id"] = origin_id
                            except Exception:
                                pass

                        print(
                            f"[umarise] ✓ Anchored {os.path.basename(local_path)} → {origin_id}",
                            file=sys.stderr,
                        )
        except Exception as e:
            print(f"[umarise] ⚠ auto_anchor error: {e}", file=sys.stderr)

        return result

    wandb.Run.log_artifact = _patched_log_artifact
    _enabled = True


def disable():
    """Disable automatic anchoring."""
    global _original_log_artifact, _enabled

    if _original_log_artifact:
        import wandb
        wandb.Run.log_artifact = _original_log_artifact
        _original_log_artifact = None
    _enabled = False
