"""
Pipeline Manager

Advanced pipeline management with versioning, experiment tracking,
and automated deployment capabilities.
"""

import pandas as pd
import numpy as np
from typing import Dict, List, Any, Optional, Tuple
import json
import os
import shutil
from datetime import datetime, timedelta
import hashlib
import joblib
from pathlib import Path
import sqlite3
from dataclasses import dataclass, asdict
import warnings


@dataclass
class ExperimentRecord:
    """Record of a pipeline experiment"""
    experiment_id: str
    pipeline_id: str
    name: str
    description: str
    created_at: str
    completed_at: Optional[str]
    status: str  # 'running', 'completed', 'failed', 'cancelled'
    config: Dict[str, Any]
    metrics: Dict[str, float]
    artifacts: Dict[str, str]  # File paths
    tags: List[str]
    parent_experiment: Optional[str] = None


@dataclass
class DeploymentRecord:
    """Record of a pipeline deployment"""
    deployment_id: str
    pipeline_id: str
    version: str
    environment: str  # 'development', 'staging', 'production'
    deployed_at: str
    status: str  # 'active', 'inactive', 'failed'
    endpoint_url: Optional[str]
    config: Dict[str, Any]
    performance_metrics: Dict[str, float]


class PipelineManager:
    """
    Advanced pipeline management system with versioning, experiment tracking,
    deployment management, and comprehensive monitoring.
    """
    
    def __init__(self, workspace_dir: str = "ai_ml_workspace"):
        """
        Initialize pipeline manager
        
        Args:
            workspace_dir: Directory for storing pipelines and experiments
        """
        self.workspace_dir = Path(workspace_dir)
        self.pipelines_dir = self.workspace_dir / "pipelines"
        self.experiments_dir = self.workspace_dir / "experiments"
        self.deployments_dir = self.workspace_dir / "deployments"
        self.models_dir = self.workspace_dir / "models"
        self.data_dir = self.workspace_dir / "data"
        
        # Create directories
        for dir_path in [self.pipelines_dir, self.experiments_dir, 
                        self.deployments_dir, self.models_dir, self.data_dir]:
            dir_path.mkdir(parents=True, exist_ok=True)
        
        # Initialize database
        self.db_path = self.workspace_dir / "pipeline_manager.db"
        self._init_database()
        
        # Active pipelines cache
        self.active_pipelines = {}
        self.experiments = {}
        self.deployments = {}
    
    def _init_database(self):
        """Initialize SQLite database for tracking"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # Create experiments table
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS experiments (
                experiment_id TEXT PRIMARY KEY,
                pipeline_id TEXT,
                name TEXT,
                description TEXT,
                created_at TEXT,
                completed_at TEXT,
                status TEXT,
                config TEXT,
                metrics TEXT,
                artifacts TEXT,
                tags TEXT,
                parent_experiment TEXT
            )
        ''')
        
        # Create deployments table
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS deployments (
                deployment_id TEXT PRIMARY KEY,
                pipeline_id TEXT,
                version TEXT,
                environment TEXT,
                deployed_at TEXT,
                status TEXT,
                endpoint_url TEXT,
                config TEXT,
                performance_metrics TEXT
            )
        ''')
        
        # Create pipelines table
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS pipelines (
                pipeline_id TEXT PRIMARY KEY,
                name TEXT,
                version TEXT,
                created_at TEXT,
                config TEXT,
                metrics TEXT,
                file_path TEXT
            )
        ''')
        
        conn.commit()
        conn.close()
    
    def register_pipeline(self, pipeline_id: str, name: str, config: Dict[str, Any],
                        pipeline_object: Any, metrics: Optional[Dict[str, float]] = None) -> str:
        """
        Register a pipeline in the system
        
        Args:
            pipeline_id: Unique pipeline identifier
            name: Pipeline name
            config: Pipeline configuration
            pipeline_object: The actual pipeline object
            metrics: Performance metrics
            
        Returns:
            Registered pipeline ID
        """
        # Save pipeline to file
        pipeline_path = self.pipelines_dir / f"{pipeline_id}.pkl"
        joblib.dump(pipeline_object, pipeline_path)
        
        # Save configuration
        config_path = self.pipelines_dir / f"{pipeline_id}_config.json"
        with open(config_path, 'w') as f:
            json.dump(config, f, indent=2)
        
        # Update database
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            INSERT OR REPLACE INTO pipelines 
            (pipeline_id, name, version, created_at, config, metrics, file_path)
            VALUES (?, ?, ?, ?, ?, ?, ?)
        ''', (
            pipeline_id,
            name,
            "1.0.0",
            datetime.now().isoformat(),
            json.dumps(config),
            json.dumps(metrics or {}),
            str(pipeline_path)
        ))
        
        conn.commit()
        conn.close()
        
        # Cache pipeline
        self.active_pipelines[pipeline_id] = {
            'object': pipeline_object,
            'config': config,
            'metrics': metrics,
            'name': name,
            'path': pipeline_path
        }
        
        return pipeline_id
    
    def load_pipeline(self, pipeline_id: str) -> Any:
        """Load a pipeline from storage"""
        if pipeline_id in self.active_pipelines:
            return self.active_pipelines[pipeline_id]['object']
        
        pipeline_path = self.pipelines_dir / f"{pipeline_id}.pkl"
        if not pipeline_path.exists():
            raise FileNotFoundError(f"Pipeline {pipeline_id} not found")
        
        pipeline = joblib.load(pipeline_path)
        
        # Load config
        config_path = self.pipelines_dir / f"{pipeline_id}_config.json"
        config = {}
        if config_path.exists():
            with open(config_path, 'r') as f:
                config = json.load(f)
        
        # Cache pipeline
        self.active_pipelines[pipeline_id] = {
            'object': pipeline,
            'config': config,
            'path': pipeline_path
        }
        
        return pipeline
    
    def create_experiment(self, name: str, description: str, pipeline_id: str,
                          config: Dict[str, Any], tags: List[str] = None,
                          parent_experiment: Optional[str] = None) -> str:
        """
        Create a new experiment
        
        Args:
            name: Experiment name
            description: Experiment description
            pipeline_id: Pipeline ID to use
            config: Experiment configuration
            tags: Experiment tags
            parent_experiment: Parent experiment ID
            
        Returns:
            Experiment ID
        """
        experiment_id = hashlib.md5(
            f"{name}_{pipeline_id}_{datetime.now().isoformat()}".encode()
        ).hexdigest()[:16]
        
        experiment = ExperimentRecord(
            experiment_id=experiment_id,
            pipeline_id=pipeline_id,
            name=name,
            description=description,
            created_at=datetime.now().isoformat(),
            completed_at=None,
            status='running',
            config=config,
            metrics={},
            artifacts={},
            tags=tags or [],
            parent_experiment=parent_experiment
        )
        
        # Create experiment directory
        experiment_dir = self.experiments_dir / experiment_id
        experiment_dir.mkdir(exist_ok=True)
        
        # Save experiment config
        with open(experiment_dir / "config.json", 'w') as f:
            json.dump(config, f, indent=2)
        
        # Update database
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            INSERT INTO experiments 
            (experiment_id, pipeline_id, name, description, created_at, status, 
             config, metrics, artifacts, tags, parent_experiment)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        ''', (
            experiment_id, pipeline_id, name, description,
            experiment.created_at, experiment.status,
            json.dumps(config), json.dumps({}), json.dumps({}),
            json.dumps(tags or []), parent_experiment
        ))
        
        conn.commit()
        conn.close()
        
        # Cache experiment
        self.experiments[experiment_id] = experiment
        
        return experiment_id
    
    def complete_experiment(self, experiment_id: str, metrics: Dict[str, float],
                           artifacts: Optional[Dict[str, str]] = None):
        """Mark an experiment as completed with results"""
        if experiment_id not in self.experiments:
            raise ValueError(f"Experiment {experiment_id} not found")
        
        experiment = self.experiments[experiment_id]
        experiment.completed_at = datetime.now().isoformat()
        experiment.status = 'completed'
        experiment.metrics = metrics
        experiment.artifacts = artifacts or {}
        
        # Save artifacts
        experiment_dir = self.experiments_dir / experiment_id
        for artifact_name, artifact_path in artifacts.items():
            if os.path.exists(artifact_path):
                shutil.copy2(artifact_path, experiment_dir / artifact_name)
        
        # Update database
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            UPDATE experiments 
            SET completed_at = ?, status = ?, metrics = ?, artifacts = ?
            WHERE experiment_id = ?
        ''', (
            experiment.completed_at, experiment.status,
            json.dumps(metrics), json.dumps(artifacts or {}),
            experiment_id
        ))
        
        conn.commit()
        conn.close()
    
    def list_experiments(self, status: Optional[str] = None, 
                         tags: Optional[List[str]] = None) -> List[Dict[str, Any]]:
        """List experiments with optional filtering"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        query = "SELECT * FROM experiments"
        params = []
        
        if status or tags:
            conditions = []
            if status:
                conditions.append("status = ?")
                params.append(status)
            
            if tags:
                for tag in tags:
                    conditions.append("tags LIKE ?")
                    params.append(f'%"{tag}"%')
            
            if conditions:
                query += " WHERE " + " AND ".join(conditions)
        
        query += " ORDER BY created_at DESC"
        
        cursor.execute(query, params)
        rows = cursor.fetchall()
        
        # Get column names
        column_names = [description[0] for description in cursor.description]
        
        experiments = []
        for row in rows:
            experiment_dict = dict(zip(column_names, row))
            # Parse JSON fields
            for field in ['config', 'metrics', 'artifacts', 'tags']:
                if experiment_dict[field]:
                    experiment_dict[field] = json.loads(experiment_dict[field])
            experiments.append(experiment_dict)
        
        conn.close()
        return experiments
    
    def get_experiment(self, experiment_id: str) -> Dict[str, Any]:
        """Get detailed experiment information"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute("SELECT * FROM experiments WHERE experiment_id = ?", (experiment_id,))
        row = cursor.fetchone()
        
        if not row:
            conn.close()
            raise ValueError(f"Experiment {experiment_id} not found")
        
        column_names = [description[0] for description in cursor.description]
        experiment_dict = dict(zip(column_names, row))
        
        # Parse JSON fields
        for field in ['config', 'metrics', 'artifacts', 'tags']:
            if experiment_dict[field]:
                experiment_dict[field] = json.loads(experiment_dict[field])
        
        conn.close()
        return experiment_dict
    
    def compare_experiments(self, experiment_ids: List[str]) -> Dict[str, Any]:
        """Compare multiple experiments"""
        experiments = []
        for exp_id in experiment_ids:
            try:
                exp_data = self.get_experiment(exp_id)
                experiments.append(exp_data)
            except ValueError:
                continue
        
        if len(experiments) < 2:
            return {"error": "Need at least 2 experiments to compare"}
        
        # Extract metrics for comparison
        all_metrics = set()
        for exp in experiments:
            all_metrics.update(exp['metrics'].keys())
        
        comparison = {
            'experiments': experiments,
            'metrics_comparison': {},
            'best_experiments': {}
        }
        
        # Compare each metric
        for metric in all_metrics:
            metric_values = []
            for exp in experiments:
                metric_values.append(exp['metrics'].get(metric, 0))
            
            comparison['metrics_comparison'][metric] = {
                'values': metric_values,
                'best_index': np.argmax(metric_values) if metric_values else 0,
                'mean': np.mean(metric_values) if metric_values else 0,
                'std': np.std(metric_values) if metric_values else 0
            }
        
        # Find best experiments for different metrics
        for metric in all_metrics:
            if metric_values:
                best_idx = np.argmax(metric_values)
                comparison['best_experiments'][metric] = experiments[best_idx]['experiment_id']
        
        return comparison
    
    def deploy_pipeline(self, pipeline_id: str, environment: str = 'development',
                       version: str = "1.0.0", config: Optional[Dict[str, Any]] = None) -> str:
        """
        Deploy a pipeline to specified environment
        
        Args:
            pipeline_id: Pipeline to deploy
            environment: Target environment
            version: Deployment version
            config: Deployment configuration
            
        Returns:
            Deployment ID
        """
        deployment_id = hashlib.md5(
            f"{pipeline_id}_{environment}_{datetime.now().isoformat()}".encode()
        ).hexdigest()[:16]
        
        # Create deployment directory
        deployment_dir = self.deployments_dir / deployment_id
        deployment_dir.mkdir(exist_ok=True)
        
        # Copy pipeline to deployment directory
        pipeline_path = self.pipelines_dir / f"{pipeline_id}.pkl"
        if pipeline_path.exists():
            shutil.copy2(pipeline_path, deployment_dir / "pipeline.pkl")
        
        # Copy config
        config_path = self.pipelines_dir / f"{pipeline_id}_config.json"
        if config_path.exists():
            shutil.copy2(config_path, deployment_dir / "config.json")
        
        # Create deployment record
        deployment = DeploymentRecord(
            deployment_id=deployment_id,
            pipeline_id=pipeline_id,
            version=version,
            environment=environment,
            deployed_at=datetime.now().isoformat(),
            status='active',
            endpoint_url=None,
            config=config or {},
            performance_metrics={}
        )
        
        # Update database
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            INSERT INTO deployments 
            (deployment_id, pipeline_id, version, environment, deployed_at, status,
             endpoint_url, config, performance_metrics)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
        ''', (
            deployment_id, pipeline_id, version, environment,
            deployment.deployed_at, deployment.status,
            deployment.endpoint_url, json.dumps(config or {}),
            json.dumps({})
        ))
        
        conn.commit()
        conn.close()
        
        # Cache deployment
        self.deployments[deployment_id] = deployment
        
        return deployment_id
    
    def list_deployments(self, environment: Optional[str] = None,
                        status: Optional[str] = None) -> List[Dict[str, Any]]:
        """List deployments with optional filtering"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        query = "SELECT * FROM deployments"
        params = []
        
        if environment or status:
            conditions = []
            if environment:
                conditions.append("environment = ?")
                params.append(environment)
            
            if status:
                conditions.append("status = ?")
                params.append(status)
            
            if conditions:
                query += " WHERE " + " AND ".join(conditions)
        
        query += " ORDER BY deployed_at DESC"
        
        cursor.execute(query, params)
        rows = cursor.fetchall()
        
        column_names = [description[0] for description in cursor.description]
        deployments = []
        
        for row in rows:
            deployment_dict = dict(zip(column_names, row))
            # Parse JSON fields
            for field in ['config', 'performance_metrics']:
                if deployment_dict[field]:
                    deployment_dict[field] = json.loads(deployment_dict[field])
            deployments.append(deployment_dict)
        
        conn.close()
        return deployments
    
    def rollback_deployment(self, deployment_id: str) -> bool:
        """Rollback a deployment"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # Update deployment status
        cursor.execute('''
            UPDATE deployments SET status = ? WHERE deployment_id = ?
        ''', ('inactive', deployment_id))
        
        conn.commit()
        conn.close()
        
        # Update cache
        if deployment_id in self.deployments:
            self.deployments[deployment_id].status = 'inactive'
        
        return True
    
    def get_pipeline_performance_history(self, pipeline_id: str) -> Dict[str, Any]:
        """Get performance history of a pipeline across experiments"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            SELECT experiment_id, name, created_at, completed_at, status, metrics
            FROM experiments 
            WHERE pipeline_id = ? AND status = 'completed'
            ORDER BY created_at ASC
        ''', (pipeline_id,))
        
        rows = cursor.fetchall()
        conn.close()
        
        history = []
        for row in rows:
            experiment_data = {
                'experiment_id': row[0],
                'name': row[1],
                'created_at': row[2],
                'completed_at': row[3],
                'status': row[4],
                'metrics': json.loads(row[5]) if row[5] else {}
            }
            history.append(experiment_data)
        
        return {
            'pipeline_id': pipeline_id,
            'history': history,
            'total_experiments': len(history)
        }
    
    def cleanup_old_experiments(self, days_old: int = 30, keep_best: bool = True) -> int:
        """Clean up old experiments"""
        cutoff_date = datetime.now() - timedelta(days=days_old)
        
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # Get experiments to delete
        cursor.execute('''
            SELECT experiment_id FROM experiments 
            WHERE created_at < ? AND status = 'completed'
        ''', (cutoff_date.isoformat(),))
        
        old_experiments = [row[0] for row in cursor.fetchall()]
        
        if keep_best and old_experiments:
            # Keep the best performing experiment(s)
            cursor.execute('''
                SELECT experiment_id, metrics FROM experiments 
                WHERE created_at < ? AND status = 'completed'
                ORDER BY created_at DESC
            ''', (cutoff_date.isoformat(),))
            
            experiments_with_metrics = cursor.fetchall()
            
            # Find best experiments based on metrics (simplified)
            best_experiments = set()
            for exp_id, metrics_json in experiments_with_metrics:
                metrics = json.loads(metrics_json) if metrics_json else {}
                # This is simplified - you'd want to be more sophisticated
                if any(metric in ['accuracy', 'r2_score'] for metric in metrics.keys()):
                    best_experiments.add(exp_id)
            
            old_experiments = [exp for exp in old_experiments if exp not in best_experiments]
        
        # Delete experiments
        deleted_count = 0
        for exp_id in old_experiments:
            # Delete from database
            cursor.execute("DELETE FROM experiments WHERE experiment_id = ?", (exp_id,))
            
            # Delete files
            exp_dir = self.experiments_dir / exp_id
            if exp_dir.exists():
                shutil.rmtree(exp_dir)
            
            deleted_count += 1
        
        conn.commit()
        conn.close()
        
        return deleted_count
    
    def export_workspace(self, export_path: str) -> bool:
        """Export entire workspace to archive"""
        try:
            shutil.make_archive(export_path, 'zip', self.workspace_dir)
            return True
        except Exception as e:
            print(f"Export failed: {e}")
            return False
    
    def import_workspace(self, import_path: str) -> bool:
        """Import workspace from archive"""
        try:
            shutil.unpack_archive(import_path, self.workspace_dir.parent)
            return True
        except Exception as e:
            print(f"Import failed: {e}")
            return False
    
    def get_workspace_summary(self) -> Dict[str, Any]:
        """Get comprehensive workspace summary"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # Get counts
        cursor.execute("SELECT COUNT(*) FROM pipelines")
        pipeline_count = cursor.fetchone()[0]
        
        cursor.execute("SELECT COUNT(*) FROM experiments")
        experiment_count = cursor.fetchone()[0]
        
        cursor.execute("SELECT COUNT(*) FROM deployments")
        deployment_count = cursor.fetchone()[0]
        
        # Get recent activity
        cursor.execute('''
            SELECT name, created_at FROM experiments 
            ORDER BY created_at DESC LIMIT 5
        ''')
        recent_experiments = cursor.fetchall()
        
        cursor.execute('''
            SELECT pipeline_id, deployed_at FROM deployments 
            WHERE status = 'active' ORDER BY deployed_at DESC LIMIT 5
        ''')
        active_deployments = cursor.fetchall()
        
        conn.close()
        
        return {
            'workspace_dir': str(self.workspace_dir),
            'pipeline_count': pipeline_count,
            'experiment_count': experiment_count,
            'deployment_count': deployment_count,
            'recent_experiments': recent_experiments,
            'active_deployments': active_deployments,
            'disk_usage': self._get_disk_usage()
        }
    
    def _get_disk_usage(self) -> Dict[str, int]:
        """Get disk usage statistics"""
        total_size = 0
        dir_sizes = {}
        
        for dir_name, dir_path in [
            ('pipelines', self.pipelines_dir),
            ('experiments', self.experiments_dir),
            ('deployments', self.deployments_dir),
            ('models', self.models_dir),
            ('data', self.data_dir)
        ]:
            if dir_path.exists():
                size = sum(f.stat().st_size for f in dir_path.rglob('*') if f.is_file())
                dir_sizes[dir_name] = size
                total_size += size
            else:
                dir_sizes[dir_name] = 0
        
        dir_sizes['total'] = total_size
        return dir_sizes
