"""
MLflow ArtifactRepository plugin that transparently anchors every artifact.

Usage:
    export MLFLOW_ARTIFACT_URI="umarise+s3://my-bucket/mlruns"
    # or: mlflow server --default-artifact-root umarise+s3://my-bucket/mlruns

Every log_artifact call will:
  1. Delegate storage to the wrapped backend.
  2. Hash the artifact locally (SHA-256, bytes never leave the machine).
  3. Submit the hash to the Umarise Core API.
  4. Write umarise.* tags to the active MLflow run.
"""

import hashlib
import os
import sys
from typing import Optional

from mlflow.store.artifact.artifact_repo import ArtifactRepository

from umarise import UmariseCore


def _hash_file(path: str) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(8192), b""):
            h.update(chunk)
    return f"sha256:{h.hexdigest()}"


def _resolve_inner_uri(uri: str) -> str:
    if uri.startswith("umarise+"):
        return uri[len("umarise+"):]
    return uri


class UmariseArtifactRepository(ArtifactRepository):
    is_plugin = True

    def __init__(self, artifact_uri: str, api_key: Optional[str] = None):
        super().__init__(artifact_uri)
        self._api_key = api_key or os.environ.get("UMARISE_API_KEY")
        inner_uri = _resolve_inner_uri(artifact_uri)
        from mlflow.store.artifact.artifact_repository_registry import get_artifact_repository
        self._inner = get_artifact_repository(inner_uri)

    def log_artifact(self, local_file, artifact_path=None):
        self._inner.log_artifact(local_file, artifact_path)
        self._anchor_file(local_file)

    def log_artifacts(self, local_dir, artifact_path=None):
        self._inner.log_artifacts(local_dir, artifact_path)
        for root, _dirs, files in os.walk(local_dir):
            for fname in files:
                self._anchor_file(os.path.join(root, fname))

    def list_artifacts(self, path=None):
        return self._inner.list_artifacts(path)

    def _download_file(self, remote_file_path, local_path):
        return self._inner._download_file(remote_file_path, local_path)

    def delete_artifacts(self, artifact_path=None):
        return self._inner.delete_artifacts(artifact_path)

    def _anchor_file(self, local_path):
        try:
            client = UmariseCore(api_key=self._api_key)
            file_hash = _hash_file(local_path)
            result = client.attest(hash_value=file_hash)
            try:
                import mlflow
                run = mlflow.active_run()
                if run:
                    mlflow.set_tag("umarise.anchored", "true")
                    mlflow.set_tag("umarise.proof_status", result.get("proof_status", "pending"))
                    mlflow.log_param(
                        f"umarise_origin_{os.path.basename(local_path)}",
                        result.get("origin_id", ""),
                    )
            except Exception:
                pass
            print(f"[umarise] Anchored {os.path.basename(local_path)} -> {result.get('origin_id', '?')}", file=sys.stderr)
        except Exception as e:
            print(f"[umarise] Anchoring failed for {os.path.basename(local_path)}: {e}", file=sys.stderr)
