"""
koci Artifact Manager

Handles artifact storage, transfer between jobs, and cleanup.
"""

import shutil
import hashlib
from pathlib import Path
from dataclasses import dataclass
from typing import Optional
from datetime import datetime

from ..core.models import Job


@dataclass
class ArtifactMetadata:
    """Metadata about a stored artifact."""
    job_name: str
    paths: list[str]
    created_at: datetime
    size_bytes: int
    checksum: str


@dataclass
class ArtifactStore:
    """Information about stored artifacts for a job."""
    job_name: str
    local_path: Path
    metadata: ArtifactMetadata


class ArtifactManager:
    """
    Manages artifacts between pipeline jobs.

    Artifacts are:
    - Copied from containers after job completion
    - Stored in a local staging directory
    - Injected into dependent job containers
    """

    def __init__(self, pipeline_name: str, run_id: Optional[str] = None):
        """
        Initialize artifact manager.

        Args:
            pipeline_name: Name of the pipeline
            run_id: Unique run identifier (generated if not provided)
        """
        self.pipeline_name = pipeline_name
        self.run_id = run_id or datetime.now().strftime('%Y%m%d-%H%M%S')

        # Base directory for artifacts
        self.base_dir = Path.home() / '.koci' / 'artifacts' / pipeline_name / self.run_id
        self.base_dir.mkdir(parents=True, exist_ok=True)

        # Track stored artifacts
        self._stores: dict[str, ArtifactStore] = {}

    @property
    def artifact_dir(self) -> Path:
        """Get the base artifact directory."""
        return self.base_dir

    def get_job_artifact_dir(self, job_name: str) -> Path:
        """Get the artifact directory for a specific job."""
        job_dir = self.base_dir / job_name
        job_dir.mkdir(parents=True, exist_ok=True)
        return job_dir

    def store_artifacts(
        self,
        job: Job,
        container_runtime,
        container_id: str,
    ) -> Optional[ArtifactStore]:
        """
        Store artifacts from a completed job.

        Args:
            job: The job that produced artifacts
            container_runtime: Container runtime instance
            container_id: ID of the job's container

        Returns:
            ArtifactStore if artifacts were stored, None otherwise
        """
        if not job.artifacts or not job.artifacts.has_outputs:
            return None

        job_dir = self.get_job_artifact_dir(job.name)
        total_size = 0
        stored_paths = []

        for path in job.artifacts.paths:
            # Copy from container to local storage
            src_path = path
            if not path.startswith('/'):
                src_path = f"{job.working_directory}/{path}"

            dest_path = job_dir / Path(path).name

            success = container_runtime.copy_from_container(
                container_id,
                src_path,
                str(dest_path.parent),
            )

            if success and dest_path.exists():
                stored_paths.append(str(dest_path))
                total_size += self._get_size(dest_path)

        if not stored_paths:
            return None

        # Calculate checksum
        checksum = self._calculate_checksum(stored_paths)

        metadata = ArtifactMetadata(
            job_name=job.name,
            paths=stored_paths,
            created_at=datetime.now(),
            size_bytes=total_size,
            checksum=checksum,
        )

        store = ArtifactStore(
            job_name=job.name,
            local_path=job_dir,
            metadata=metadata,
        )

        self._stores[job.name] = store
        return store

    def inject_artifacts(
        self,
        job: Job,
        container_runtime,
        container_id: str,
    ) -> bool:
        """
        Inject consumed artifacts into a job's container.

        Args:
            job: The job that consumes artifacts
            container_runtime: Container runtime instance
            container_id: ID of the job's container

        Returns:
            True if all artifacts were injected successfully
        """
        if not job.artifacts or not job.artifacts.has_inputs:
            return True

        success = True

        for source_job_name in job.artifacts.consume:
            store = self._stores.get(source_job_name)

            if not store:
                print(f"Warning: No artifacts found from job '{source_job_name}'")
                continue

            # Copy artifacts to container
            for local_path in store.metadata.paths:
                dest_path = job.working_directory

                result = container_runtime.copy_to_container(
                    container_id,
                    local_path,
                    dest_path,
                )

                if not result:
                    print(f"Failed to inject artifact: {local_path}")
                    success = False

        return success

    def get_artifact_store(self, job_name: str) -> Optional[ArtifactStore]:
        """Get stored artifacts for a job."""
        return self._stores.get(job_name)

    def list_artifacts(self) -> list[ArtifactStore]:
        """List all stored artifacts."""
        return list(self._stores.values())

    def cleanup(self, keep_days: int = 0):
        """
        Clean up artifact storage.

        Args:
            keep_days: Keep artifacts newer than this many days (0 = delete all)
        """
        if keep_days == 0:
            # Delete everything for this run
            if self.base_dir.exists():
                shutil.rmtree(self.base_dir)
            self._stores.clear()
            return

        # Delete artifacts older than keep_days
        cutoff = datetime.now().timestamp() - (keep_days * 86400)

        for job_name, store in list(self._stores.items()):
            if store.metadata.created_at.timestamp() < cutoff:
                if store.local_path.exists():
                    shutil.rmtree(store.local_path)
                del self._stores[job_name]

    def cleanup_old_runs(self, keep_runs: int = 5):
        """
        Clean up old pipeline runs.

        Args:
            keep_runs: Number of recent runs to keep
        """
        pipeline_dir = Path.home() / '.koci' / 'artifacts' / self.pipeline_name

        if not pipeline_dir.exists():
            return

        # Get all run directories sorted by modification time
        run_dirs = sorted(
            pipeline_dir.iterdir(),
            key=lambda p: p.stat().st_mtime,
            reverse=True,
        )

        # Remove old runs
        for run_dir in run_dirs[keep_runs:]:
            if run_dir.is_dir():
                shutil.rmtree(run_dir)

    def _get_size(self, path: Path) -> int:
        """Get total size of a path (file or directory)."""
        if path.is_file():
            return path.stat().st_size

        total = 0
        for item in path.rglob('*'):
            if item.is_file():
                total += item.stat().st_size
        return total

    def _calculate_checksum(self, paths: list[str]) -> str:
        """Calculate combined checksum of all artifact paths."""
        hasher = hashlib.sha256()

        for path in sorted(paths):
            path_obj = Path(path)
            if path_obj.is_file():
                with open(path_obj, 'rb') as f:
                    for chunk in iter(lambda: f.read(8192), b''):
                        hasher.update(chunk)
            elif path_obj.is_dir():
                for file_path in sorted(path_obj.rglob('*')):
                    if file_path.is_file():
                        with open(file_path, 'rb') as f:
                            for chunk in iter(lambda: f.read(8192), b''):
                                hasher.update(chunk)

        return hasher.hexdigest()[:16]


def get_artifact_manager(pipeline_name: str, run_id: Optional[str] = None) -> ArtifactManager:
    """
    Get an artifact manager instance.

    Args:
        pipeline_name: Name of the pipeline
        run_id: Optional run identifier

    Returns:
        ArtifactManager instance
    """
    return ArtifactManager(pipeline_name, run_id)
