"""
Version Manager

Advanced version management system for ML models, datasets, and pipelines
with semantic versioning, dependency tracking, and rollback capabilities.
"""

import os
import json
import pickle
import hashlib
from typing import Dict, List, Any, Optional, Union, Tuple
from datetime import datetime
from pathlib import Path
import shutil
import sqlite3
from dataclasses import dataclass, asdict
from enum import Enum
import warnings


class VersionType(Enum):
    """Types of version changes"""
    MAJOR = "major"    # Breaking changes
    MINOR = "minor"    # New features, backward compatible
    PATCH = "patch"    # Bug fixes, backward compatible


@dataclass
class VersionInfo:
    """Version information"""
    major: int
    minor: int
    patch: int
    build: Optional[str] = None
    prerelease: Optional[str] = None
    
    def __str__(self):
        version = f"{self.major}.{self.minor}.{self.patch}"
        if self.prerelease:
            version += f"-{self.prerelease}"
        if self.build:
            version += f"+{self.build}"
        return version
    
    @classmethod
    def from_string(cls, version_str: str):
        """Create VersionInfo from string"""
        parts = version_str.split('+')
        build = parts[1] if len(parts) > 1 else None
        version_part = parts[0]
        
        parts = version_part.split('-')
        prerelease = parts[1] if len(parts) > 1 else None
        version_part = parts[0]
        
        major, minor, patch = map(int, version_part.split('.'))
        
        return cls(major=major, minor=minor, patch=patch, build=build, prerelease=prerelease)
    
    def bump(self, version_type: VersionType):
        """Bump version"""
        if version_type == VersionType.MAJOR:
            return VersionInfo(self.major + 1, 0, 0, self.build, self.prerelease)
        elif version_type == VersionType.MINOR:
            return VersionInfo(self.major, self.minor + 1, 0, self.build, self.prerelease)
        elif version_type == VersionType.PATCH:
            return VersionInfo(self.major, self.minor, self.patch + 1, self.build, self.prerelease)


@dataclass
class Artifact:
    """Model or dataset artifact"""
    artifact_id: str
    name: str
    version: VersionInfo
    artifact_type: str  # 'model', 'dataset', 'pipeline', 'config'
    file_path: str
    checksum: str
    created_at: str
    created_by: str
    metadata: Dict[str, Any]
    dependencies: List[str]  # List of artifact_ids
    tags: List[str]


class VersionManager:
    """
    Advanced version management system for ML artifacts with semantic versioning,
    dependency tracking, and comprehensive metadata management.
    """
    
    def __init__(self, storage_path: str = "version_storage"):
        """
        Initialize version manager
        
        Args:
            storage_path: Path to store versioned artifacts
        """
        self.storage_path = Path(storage_path)
        self.storage_path.mkdir(exist_ok=True)
        
        # Create subdirectories
        self.artifacts_dir = self.storage_path / "artifacts"
        self.metadata_dir = self.storage_path / "metadata"
        self.db_path = self.storage_path / "versions.db"
        
        for dir_path in [self.artifacts_dir, self.metadata_dir]:
            dir_path.mkdir(exist_ok=True)
        
        # Initialize database
        self._init_database()
        
        # Cache
        self.artifact_cache = {}
    
    def _init_database(self):
        """Initialize SQLite database for version tracking"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS artifacts (
                artifact_id TEXT PRIMARY KEY,
                name TEXT,
                version TEXT,
                artifact_type TEXT,
                file_path TEXT,
                checksum TEXT,
                created_at TEXT,
                created_by TEXT,
                metadata TEXT,
                dependencies TEXT,
                tags TEXT,
                is_latest BOOLEAN
            )
        ''')
        
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS version_history (
                history_id INTEGER PRIMARY KEY AUTOINCREMENT,
                artifact_id TEXT,
                version TEXT,
                action TEXT,
                timestamp TEXT,
                description TEXT,
                FOREIGN KEY (artifact_id) REFERENCES artifacts (artifact_id)
            )
        ''')
        
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS dependencies (
                dependency_id INTEGER PRIMARY KEY AUTOINCREMENT,
                artifact_id TEXT,
                depends_on TEXT,
                version_constraint TEXT,
                FOREIGN KEY (artifact_id) REFERENCES artifacts (artifact_id)
            )
        ''')
        
        conn.commit()
        conn.close()
    
    def _calculate_checksum(self, file_path: str) -> str:
        """Calculate SHA-256 checksum of file"""
        hash_sha256 = hashlib.sha256()
        with open(file_path, "rb") as f:
            for chunk in iter(lambda: f.read(4096), b""):
                hash_sha256.update(chunk)
        return hash_sha256.hexdigest()
    
    def _generate_artifact_id(self, name: str, version: VersionInfo) -> str:
        """Generate unique artifact ID"""
        content = f"{name}_{str(version)}_{datetime.now().isoformat()}"
        return hashlib.md5(content.encode()).hexdigest()[:16]
    
    def save_artifact(self, artifact_path: str, name: str, version: VersionInfo,
                     artifact_type: str, creator: str = "user",
                     metadata: Optional[Dict[str, Any]] = None,
                     dependencies: Optional[List[str]] = None,
                     tags: Optional[List[str]] = None) -> str:
        """
        Save an artifact with versioning
        
        Args:
            artifact_path: Path to the artifact file
            name: Name of the artifact
            version: Version information
            artifact_type: Type of artifact
            creator: Creator of the artifact
            metadata: Additional metadata
            dependencies: List of dependency artifact IDs
            tags: Tags for the artifact
            
        Returns:
            Artifact ID
        """
        if not os.path.exists(artifact_path):
            raise FileNotFoundError(f"Artifact not found: {artifact_path}")
        
        # Generate artifact ID
        artifact_id = self._generate_artifact_id(name, version)
        
        # Calculate checksum
        checksum = self._calculate_checksum(artifact_path)
        
        # Copy artifact to storage
        artifact_storage_path = self.artifacts_dir / artifact_id
        shutil.copy2(artifact_path, artifact_storage_path)
        
        # Create artifact object
        artifact = Artifact(
            artifact_id=artifact_id,
            name=name,
            version=version,
            artifact_type=artifact_type,
            file_path=str(artifact_storage_path),
            checksum=checksum,
            created_at=datetime.now().isoformat(),
            created_by=creator,
            metadata=metadata or {},
            dependencies=dependencies or [],
            tags=tags or []
        )
        
        # Save to database
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            INSERT INTO artifacts 
            (artifact_id, name, version, artifact_type, file_path, checksum, 
             created_at, created_by, metadata, dependencies, tags, is_latest)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        ''', (
            artifact.artifact_id, artifact.name, str(artifact.version),
            artifact.artifact_type, artifact.file_path, artifact.checksum,
            artifact.created_at, artifact.created_by, json.dumps(artifact.metadata),
            json.dumps(artifact.dependencies), json.dumps(artifact.tags), True
        ))
        
        # Update previous versions to not be latest
        cursor.execute('''
            UPDATE artifacts SET is_latest = FALSE 
            WHERE name = ? AND artifact_id != ? AND artifact_type = ?
        ''', (name, artifact_id, artifact_type))
        
        # Add to version history
        cursor.execute('''
            INSERT INTO version_history 
            (artifact_id, version, action, timestamp, description)
            VALUES (?, ?, ?, ?, ?)
        ''', (artifact_id, str(version), "CREATE", artifact.created_at, 
               f"Created {artifact_type} {name} version {version}"))
        
        # Add dependencies
        if dependencies:
            for dep_id in dependencies:
                cursor.execute('''
                    INSERT INTO dependencies (artifact_id, depends_on, version_constraint)
                    VALUES (?, ?, ?)
                ''', (artifact_id, dep_id, ">=0.0.0"))
        
        conn.commit()
        conn.close()
        
        # Cache artifact
        self.artifact_cache[artifact_id] = artifact
        
        print(f"Saved artifact: {name} v{version} (ID: {artifact_id})")
        return artifact_id
    
    def load_artifact(self, artifact_id: str) -> Artifact:
        """Load artifact metadata"""
        if artifact_id in self.artifact_cache:
            return self.artifact_cache[artifact_id]
        
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            SELECT * FROM artifacts WHERE artifact_id = ?
        ''', (artifact_id,))
        
        row = cursor.fetchone()
        conn.close()
        
        if not row:
            raise ValueError(f"Artifact not found: {artifact_id}")
        
        artifact = Artifact(
            artifact_id=row[0],
            name=row[1],
            version=VersionInfo.from_string(row[2]),
            artifact_type=row[3],
            file_path=row[4],
            checksum=row[5],
            created_at=row[6],
            created_by=row[7],
            metadata=json.loads(row[8]),
            dependencies=json.loads(row[9]),
            tags=json.loads(row[10])
        )
        
        self.artifact_cache[artifact_id] = artifact
        return artifact
    
    def get_artifact_file(self, artifact_id: str) -> str:
        """Get the file path for an artifact"""
        artifact = self.load_artifact(artifact_id)
        
        if not os.path.exists(artifact.file_path):
            raise FileNotFoundError(f"Artifact file not found: {artifact.file_path}")
        
        return artifact.file_path
    
    def list_artifacts(self, name: Optional[str] = None, 
                      artifact_type: Optional[str] = None,
                      tags: Optional[List[str]] = None,
                      latest_only: bool = False) -> List[Artifact]:
        """List artifacts with optional filtering"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        query = "SELECT * FROM artifacts WHERE 1=1"
        params = []
        
        if name:
            query += " AND name = ?"
            params.append(name)
        
        if artifact_type:
            query += " AND artifact_type = ?"
            params.append(artifact_type)
        
        if latest_only:
            query += " AND is_latest = TRUE"
        
        query += " ORDER BY created_at DESC"
        
        cursor.execute(query, params)
        rows = cursor.fetchall()
        conn.close()
        
        artifacts = []
        for row in rows:
            artifact_tags = json.loads(row[10])
            
            # Filter by tags if specified
            if tags and not any(tag in artifact_tags for tag in tags):
                continue
            
            artifact = Artifact(
                artifact_id=row[0],
                name=row[1],
                version=VersionInfo.from_string(row[2]),
                artifact_type=row[3],
                file_path=row[4],
                checksum=row[5],
                created_at=row[6],
                created_by=row[7],
                metadata=json.loads(row[8]),
                dependencies=json.loads(row[9]),
                tags=artifact_tags
            )
            artifacts.append(artifact)
        
        return artifacts
    
    def get_latest_version(self, name: str, artifact_type: Optional[str] = None) -> Optional[Artifact]:
        """Get the latest version of an artifact"""
        artifacts = self.list_artifacts(name=name, artifact_type=artifact_type, latest_only=True)
        return artifacts[0] if artifacts else None
    
    def get_version_history(self, name: str, artifact_type: Optional[str] = None) -> List[Artifact]:
        """Get version history for an artifact"""
        return self.list_artifacts(name=name, artifact_type=artifact_type)
    
    def create_new_version(self, name: str, version_type: VersionType,
                          artifact_type: str, artifact_path: str,
                          creator: str = "user",
                          metadata: Optional[Dict[str, Any]] = None,
                          tags: Optional[List[str]] = None) -> str:
        """Create a new version of an existing artifact"""
        # Get latest version
        latest = self.get_latest_version(name, artifact_type)
        
        if latest:
            new_version = latest.version.bump(version_type)
        else:
            new_version = VersionInfo(1, 0, 0)
        
        return self.save_artifact(
            artifact_path=artifact_path,
            name=name,
            version=new_version,
            artifact_type=artifact_type,
            creator=creator,
            metadata=metadata,
            tags=tags
        )
    
    def rollback_to_version(self, name: str, version: VersionInfo, 
                           artifact_type: str) -> str:
        """Rollback to a specific version"""
        # Find the specific version
        artifacts = self.list_artifacts(name=name, artifact_type=artifact_type)
        
        target_artifact = None
        for artifact in artifacts:
            if artifact.version == version:
                target_artifact = artifact
                break
        
        if not target_artifact:
            raise ValueError(f"Version {version} not found for {name}")
        
        # Create new version with rollback content
        rollback_version = target_artifact.version.bump(VersionType.MINOR)
        rollback_version.prerelease = "rollback"
        
        return self.save_artifact(
            artifact_path=target_artifact.file_path,
            name=f"{name}_rollback",
            version=rollback_version,
            artifact_type=artifact_type,
            creator="system",
            metadata={
                "rollback_from": str(target_artifact.version),
                "rollback_reason": "Manual rollback"
            },
            tags=["rollback"]
        )
    
    def compare_versions(self, artifact_id1: str, artifact_id2: str) -> Dict[str, Any]:
        """Compare two versions of an artifact"""
        artifact1 = self.load_artifact(artifact_id1)
        artifact2 = self.load_artifact(artifact_id2)
        
        comparison = {
            "artifact1": {
                "id": artifact1.artifact_id,
                "name": artifact1.name,
                "version": str(artifact1.version),
                "created_at": artifact1.created_at
            },
            "artifact2": {
                "id": artifact2.artifact_id,
                "name": artifact2.name,
                "version": str(artifact2.version),
                "created_at": artifact2.created_at
            },
            "differences": {}
        }
        
        # Compare versions
        if artifact1.version > artifact2.version:
            comparison["differences"]["version"] = f"{artifact1.version} is newer than {artifact2.version}"
        elif artifact1.version < artifact2.version:
            comparison["differences"]["version"] = f"{artifact2.version} is newer than {artifact1.version}"
        
        # Compare checksums
        if artifact1.checksum != artifact2.checksum:
            comparison["differences"]["content"] = "Artifacts have different content"
        
        # Compare metadata
        metadata_diff = set(artifact1.metadata.items()) ^ set(artifact2.metadata.items())
        if metadata_diff:
            comparison["differences"]["metadata"] = list(metadata_diff)
        
        # Compare dependencies
        deps_diff = set(artifact1.dependencies) ^ set(artifact2.dependencies)
        if deps_diff:
            comparison["differences"]["dependencies"] = list(deps_diff)
        
        return comparison
    
    def get_dependency_graph(self, artifact_id: str) -> Dict[str, Any]:
        """Get dependency graph for an artifact"""
        artifact = self.load_artifact(artifact_id)
        
        graph = {
            "artifact_id": artifact_id,
            "name": artifact.name,
            "version": str(artifact.version),
            "dependencies": [],
            "dependents": []
        }
        
        # Get dependencies
        for dep_id in artifact.dependencies:
            try:
                dep_artifact = self.load_artifact(dep_id)
                graph["dependencies"].append({
                    "artifact_id": dep_id,
                    "name": dep_artifact.name,
                    "version": str(dep_artifact.version)
                })
            except ValueError:
                graph["dependencies"].append({
                    "artifact_id": dep_id,
                    "name": "Unknown",
                    "version": "Unknown"
                })
        
        # Get dependents (artifacts that depend on this one)
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            SELECT DISTINCT artifact_id FROM dependencies WHERE depends_on = ?
        ''', (artifact_id,))
        
        dependent_rows = cursor.fetchall()
        conn.close()
        
        for row in dependent_rows:
            try:
                dependent = self.load_artifact(row[0])
                graph["dependents"].append({
                    "artifact_id": row[0],
                    "name": dependent.name,
                    "version": str(dependent.version)
                })
            except ValueError:
                graph["dependents"].append({
                    "artifact_id": row[0],
                    "name": "Unknown",
                    "version": "Unknown"
                })
        
        return graph
    
    def validate_dependencies(self, artifact_id: str) -> Dict[str, Any]:
        """Validate dependencies for an artifact"""
        artifact = self.load_artifact(artifact_id)
        
        validation = {
            "artifact_id": artifact_id,
            "valid": True,
            "issues": [],
            "missing_dependencies": []
        }
        
        for dep_id in artifact.dependencies:
            try:
                dep_artifact = self.load_artifact(dep_id)
                
                # Check if file exists
                if not os.path.exists(dep_artifact.file_path):
                    validation["issues"].append(f"Dependency file missing: {dep_id}")
                    validation["valid"] = False
                    validation["missing_dependencies"].append(dep_id)
                
            except ValueError:
                validation["issues"].append(f"Dependency not found: {dep_id}")
                validation["valid"] = False
                validation["missing_dependencies"].append(dep_id)
        
        return validation
    
    def cleanup_old_versions(self, name: str, artifact_type: str, 
                           keep_count: int = 5) -> List[str]:
        """Clean up old versions, keeping only the latest N"""
        artifacts = self.get_version_history(name, artifact_type)
        
        if len(artifacts) <= keep_count:
            return []
        
        # Sort by version (newest first) and keep the latest N
        artifacts.sort(key=lambda x: x.version, reverse=True)
        to_keep = artifacts[:keep_count]
        to_remove = artifacts[keep_count:]
        
        removed_ids = []
        
        for artifact in to_remove:
            # Remove from database
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            cursor.execute('DELETE FROM artifacts WHERE artifact_id = ?', (artifact.artifact_id,))
            cursor.execute('DELETE FROM dependencies WHERE artifact_id = ?', (artifact.artifact_id,))
            cursor.execute('DELETE FROM dependencies WHERE depends_on = ?', (artifact.artifact_id,))
            
            conn.commit()
            conn.close()
            
            # Remove file
            if os.path.exists(artifact.file_path):
                os.remove(artifact.file_path)
            
            # Remove from cache
            if artifact.artifact_id in self.artifact_cache:
                del self.artifact_cache[artifact.artifact_id]
            
            removed_ids.append(artifact.artifact_id)
        
        print(f"Cleaned up {len(removed_ids)} old versions of {name}")
        return removed_ids
    
    def export_artifact(self, artifact_id: str, export_path: str):
        """Export an artifact to a specified path"""
        artifact = self.load_artifact(artifact_id)
        
        # Copy artifact file
        shutil.copy2(artifact.file_path, export_path)
        
        # Create metadata file
        metadata_path = export_path + ".metadata.json"
        with open(metadata_path, 'w') as f:
            json.dump(asdict(artifact), f, indent=2, default=str)
        
        print(f"Exported artifact {artifact_id} to {export_path}")
    
    def import_artifact(self, artifact_path: str, metadata_path: str) -> str:
        """Import an artifact from exported files"""
        # Load metadata
        with open(metadata_path, 'r') as f:
            metadata = json.load(f)
        
        # Reconstruct artifact object
        artifact = Artifact(**metadata)
        
        # Save to version manager
        return self.save_artifact(
            artifact_path=artifact_path,
            name=artifact.name,
            version=artifact.version,
            artifact_type=artifact.artifact_type,
            creator=artifact.created_by,
            metadata=artifact.metadata,
            dependencies=artifact.dependencies,
            tags=artifact.tags
        )
    
    def get_statistics(self) -> Dict[str, Any]:
        """Get version manager statistics"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # Total artifacts
        cursor.execute("SELECT COUNT(*) FROM artifacts")
        total_artifacts = cursor.fetchone()[0]
        
        # Artifacts by type
        cursor.execute("SELECT artifact_type, COUNT(*) FROM artifacts GROUP BY artifact_type")
        artifacts_by_type = dict(cursor.fetchall())
        
        # Latest versions
        cursor.execute("SELECT COUNT(*) FROM artifacts WHERE is_latest = TRUE")
        latest_versions = cursor.fetchone()[0]
        
        # Total dependencies
        cursor.execute("SELECT COUNT(*) FROM dependencies")
        total_dependencies = cursor.fetchone()[0]
        
        conn.close()
        
        return {
            "total_artifacts": total_artifacts,
            "latest_versions": latest_versions,
            "artifacts_by_type": artifacts_by_type,
            "total_dependencies": total_dependencies,
            "storage_path": str(self.storage_path)
        }
