"""MLflow and DVC Object Storage Integration.

Triggers when MLflow is present alongside MinIO or S3.
Wires the tracking server's artifact storage correctly, and if DVC
is also active, configures the DVC remote backend.
"""

from __future__ import annotations

from nikame.codegen.integrations.base import BaseIntegration


class MLflowStorageIntegration(BaseIntegration):
    """Generates ML tracking storage configurations."""

    # Dynamic trigger logic handles this
    REQUIRED_MODULES = []

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.has_dvc = "dvc" in self.active_modules
        self.storage_backend = "minio" if "minio" in self.active_modules else "s3"

    @classmethod
    def should_trigger(cls, active_modules: set[str], active_features: set[str]) -> bool:
        """Custom trigger: Requires MLflow AND (MinIO or AWS S3)."""
        has_mlflow = "mlflow" in active_modules
        has_storage = any(m in active_modules for m in ["minio", "s3"])
        return has_mlflow and has_storage

    def generate_core(self) -> list[tuple[str, str]]:
        # The frontend/FastAPI logic is mostly handled directly inside the API module
        # but here we generate a robust DVC/MLflow setup script for the infrastructure.
        scripts = self._generate_setup_script()
        return [("infra/scripts/setup_ml_storage.sh", scripts)]

    def generate_lifespan(self) -> str:
        # FastAPI tracking lifespan is handled in the FastAPIModule, 
        # but we could add safety checks here if necessary.
        return ""

    def generate_health(self) -> dict[str, str]:
        return {}

    def generate_metrics(self) -> str:
        return ""

    def generate_guide(self) -> str:
        dvc_note = ""
        if self.has_dvc:
            dvc_note = f"\n**Data Version Control (DVC)** is also bound to this {self.storage_backend} instance. Your `dvc push` commands will automatically route to the `dvc` bucket."
            
        return f"""
### MLflow Artifact Storage Setup
**Status:** Active 🟢 
**Artifact Backend:** `{self.storage_backend}`

MLflow is natively tracking metadata in Postgres and resolving artifact payloads (models, plots) against the locally deployed Object Storage. {dvc_note}
"""

    def _generate_setup_script(self) -> str:
        script = f"""#!/bin/bash
# DVC and MLflow Storage Automated Initialization
# Generated by NIKAME Matrix Engine

echo "Setting up ML storage buckets in {self.storage_backend}..."

"""
        if self.has_dvc:
            script += """
# Configure DVC Remote
if [ -d ".dvc" ]; then
    echo "Configuring DVC remote..."
    dvc remote add -d default s3://dvc --force
    dvc remote modify default endpointurl http://localhost:9000
    dvc remote modify default access_key_id admin
    dvc remote modify default secret_access_key admin123
else
    echo "DVC not initialized in this directory. Run 'dvc init' first."
fi
"""
        return script
