"""
MlvaultArtifactRepository — MLflow artifact store plugin backed by Jackal decentralized storage.

URI format:  mlvault://<project-name>
MLflow appends  /<mlflow_run_id>/artifacts  to form the full artifact_uri.

Staging model (like Git):
  log_artifact  -> instant local copy  (~/.mlvault/mlflow-staging/<run_id>/artifacts/)
  mlvault commit <run_id> -> bundle + encrypt + upload to Jackal  (run after training)

Download path:
  _download_file -> delegates to mlvault storage layer (requires Jackal env vars)
"""

import os
import shutil
import sys
from pathlib import Path
from typing import Optional
from urllib.parse import urlparse

from mlflow.store.artifact.artifact_repo import ArtifactRepository
from mlflow.entities import FileInfo

from mlvault_mlflow import staging

# mlvault core lives one directory up (same monorepo).
_MLVAULT_CORE = Path(__file__).parent.parent.parent / "mlvault"
if str(_MLVAULT_CORE.parent) not in sys.path:
    sys.path.insert(0, str(_MLVAULT_CORE.parent))


def _parse_uri(artifact_uri: str) -> tuple[str, str]:
    """
    Parse  mlvault://project-name/<mlflow_run_id>/artifacts[/sub/path]
    Returns (project, mlflow_run_id).
    MLflow sets artifact_uri to mlvault://project/<run_id>/artifacts automatically.
    """
    parsed = urlparse(artifact_uri)
    project = parsed.netloc  # project-name
    # path is like /<run_id>/artifacts or /<run_id>/artifacts/subdir
    parts = [p for p in parsed.path.split("/") if p]
    run_id = parts[0] if parts else "unknown"
    return project, run_id


class MlvaultArtifactRepository(ArtifactRepository):
    """MLflow artifact repository backed by Jackal decentralised storage."""

    def __init__(self, artifact_uri: str, **kwargs):
        super().__init__(artifact_uri, **kwargs)
        self.project, self.mlflow_run_id = _parse_uri(artifact_uri)

    # ── Write path (staging) ──────────────────────────────────────────────────

    def log_artifact(self, local_file: str, artifact_path: Optional[str] = None) -> None:
        """Stage a single file locally (instant — no Jackal I/O)."""
        staging.write_meta(self.mlflow_run_id, self.project)
        staging.stage_file(self.mlflow_run_id, local_file, artifact_path)

    def log_artifacts(self, local_dir: str, artifact_path: Optional[str] = None) -> None:
        """Stage all files under local_dir."""
        staging.write_meta(self.mlflow_run_id, self.project)
        src = Path(local_dir)
        for f in src.rglob("*"):
            if f.is_file():
                rel = f.relative_to(src).parent
                sub = str(artifact_path / rel) if artifact_path else str(rel) if str(rel) != "." else None
                staging.stage_file(self.mlflow_run_id, str(f), sub or artifact_path)

    # ── List ──────────────────────────────────────────────────────────────────

    def list_artifacts(self, path: Optional[str] = None) -> list[FileInfo]:
        """List staged artifacts. Falls back to Jackal if a committed manifest is present."""
        staged = staging.list_staged_files(self.mlflow_run_id)
        base = Path.home() / ".mlvault" / "mlflow-staging" / self.mlflow_run_id / "artifacts"

        results = []
        for f in staged:
            rel = f.relative_to(base)
            # Filter to path prefix if requested
            if path and not str(rel).startswith(path):
                continue
            results.append(FileInfo(
                path=str(rel).replace("\\", "/"),
                is_dir=False,
                file_size=f.stat().st_size,
            ))
        return results

    # ── Read path (download from Jackal) ─────────────────────────────────────

    def _download_file(self, remote_file_path: str, local_path: str) -> None:
        """
        Download a single artifact.

        First checks the local staging area (fast path for within-session access).
        Falls back to downloading from Jackal using the committed CID.
        """
        staged_file = (
            Path.home() / ".mlvault" / "mlflow-staging"
            / self.mlflow_run_id / "artifacts" / remote_file_path
        )
        if staged_file.exists():
            shutil.copy2(staged_file, local_path)
            return

        # Jackal download path — requires committed manifest with jackal_path
        from mlvault import snapshot, storage, crypto
        import tempfile

        try:
            manifest = snapshot.load_manifest(self.mlflow_run_id)
        except FileNotFoundError:
            raise FileNotFoundError(
                f"Artifact '{remote_file_path}' not found in staging and no committed "
                f"manifest exists for run {self.mlflow_run_id}. "
                f"Run 'mlvault commit {self.mlflow_run_id}' first."
            )

        cid = manifest.get("jackal_path")
        if not cid:
            raise RuntimeError(f"Run {self.mlflow_run_id} has not been uploaded to Jackal.")

        with tempfile.TemporaryDirectory() as tmp:
            enc_path = Path(tmp) / "bundle.enc"
            storage.download(cid, enc_path)
            plaintext = crypto.decrypt_bundle(enc_path, self.mlflow_run_id)

            import tarfile, io
            with tarfile.open(fileobj=io.BytesIO(plaintext), mode="r:gz") as tar:
                # Find the member matching remote_file_path
                target = None
                for member in tar.getmembers():
                    if member.name.endswith(remote_file_path.replace("\\", "/")):
                        target = member
                        break
                if not target:
                    raise FileNotFoundError(
                        f"'{remote_file_path}' not found in Jackal bundle for {self.mlflow_run_id}"
                    )
                Path(local_path).parent.mkdir(parents=True, exist_ok=True)
                with tar.extractfile(target) as src_f, open(local_path, "wb") as dst_f:
                    dst_f.write(src_f.read())
