import hashlib, os, sys
from typing import Optional
from umarise_wandb.queue import AnchorQueue, AnchorJob
from umarise_wandb.ledger import ReconciliationLedger

_queue = None
_ledger = None

def _get_queue():
    global _queue
    if _queue is None:
        _queue = AnchorQueue()
        _queue.on_success = _on_success
        _queue.on_failure = _on_failure
        _queue.start()
    return _queue

def _get_ledger():
    global _ledger
    if _ledger is None: _ledger = ReconciliationLedger()
    return _ledger

def _on_success(job, origin_id, proof_status):
    _get_ledger().record(source="wandb", file_path=job.file_path, file_hash=job.file_hash, origin_id=origin_id, proof_status=proof_status)
    ctx = job.context or {}
    run = ctx.get("run")
    if run:
        try:
            run.summary["umarise_anchored"] = True
            run.summary["umarise_origin_id"] = origin_id
        except Exception: pass

def _on_failure(job, error):
    _get_ledger().record_failure(source="wandb", file_path=job.file_path, file_hash=job.file_hash, error=error)

def _hash_file(path):
    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=None):
    from umarise import UmariseCore
    key = api_key or os.environ.get("UMARISE_API_KEY")
    if not key: raise RuntimeError("UMARISE_API_KEY not set.")
    return UmariseCore(api_key=key)

def anchor_file(file_path, run=None, api_key=None, sync=False):
    try:
        file_hash = _hash_file(file_path)
        if sync:
            client = _get_client(api_key)
            result = client.attest(f"sha256:{file_hash}")
            origin_id = getattr(result, "origin_id", None) or "unknown"
            _get_ledger().record(source="wandb", file_path=file_path, file_hash=file_hash, origin_id=origin_id)
            if run:
                try:
                    run.summary["umarise_anchored"] = True
                    run.summary["umarise_origin_id"] = origin_id
                except Exception: pass
            print(f"[umarise] ✓ {os.path.basename(file_path)} → {origin_id}", file=sys.stderr)
            return origin_id
        else:
            job = AnchorJob(file_path=file_path, file_hash=file_hash, api_key=api_key or os.environ.get("UMARISE_API_KEY"), context={"run": run})
            _get_queue().submit(job)
            return None
    except Exception as e:
        print(f"[umarise] ⚠ {file_path}: {e}", file=sys.stderr)
        return None

def anchor_artifact(artifact, run=None, api_key=None, sync=False):
    results = []
    try:
        for entry in artifact.manifest.entries.values():
            lp = getattr(entry, "local_path", None)
            if lp and os.path.isfile(lp):
                oid = anchor_file(lp, run=run, api_key=api_key, sync=sync)
                results.append((lp, oid))
    except Exception: pass
    return results

def anchor_logged_artifact(artifact_name, version="latest", api_key=None, sync=False):
    import wandb
    if not wandb.run: raise RuntimeError("No active W&B run.")
    ref = f"{artifact_name}:{version}" if ":" not in artifact_name else artifact_name
    artifact = wandb.run.use_artifact(ref)
    d = artifact.download()
    results = []
    for root, _, files in os.walk(d):
        for f in files:
            fp = os.path.join(root, f)
            oid = anchor_file(fp, run=wandb.run, api_key=api_key, sync=sync)
            results.append((fp, oid))
    return results

def flush(timeout=10.0):
    if _queue: _queue.shutdown(timeout=timeout)
def pending_anchors():
    return _queue.pending if _queue else 0
def ledger():
    return _get_ledger()
