import sys
import io
import time
import trace
import traceback
import psutil
import pandas as pd
import numpy as np
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score

# Default healthy pipeline code string
DEFAULT_PIPELINE_CODE = """import numpy as np
import pandas as pd
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score

def run_ml_pipeline():
    print("[Pipeline] Ingesting features and targets...")
    X_raw, y_raw = make_classification(
        n_samples=1200, n_features=10, n_informative=8,
        n_redundant=2, random_state=42
    )
    
    df = pd.DataFrame(X_raw, columns=[f"feat_{i}" for i in range(10)])
    df["target"] = y_raw
    
    print("[Pipeline] Preprocessing data and handling nulls...")
    # Clean data baseline
    df = df.dropna()
    
    X = df.drop(columns=["target"])
    y = df["target"]
    
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
    
    print("[Pipeline] Training Random Forest model...")
    model = RandomForestClassifier(n_estimators=50, max_depth=6, random_state=42)
    model.fit(X_train, y_train)
    
    preds = model.predict(X_test)
    acc = accuracy_score(y_test, preds)
    loss = float(1.0 - acc)
    
    print(f"[Pipeline Execution Finished] Accuracy: {acc:.4f}, Loss: {loss:.4f}")
    return {"accuracy": float(acc), "loss": float(loss), "samples": len(df)}

result = run_ml_pipeline()
"""

# Fault scenario templates to inject into pipeline
FAULT_TEMPLATES = {
    "DATA_DRIFT": """import numpy as np
import pandas as pd
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score

def run_ml_pipeline():
    print("[Pipeline] Ingesting features and targets...")
    X_raw, y_raw = make_classification(n_samples=1200, n_features=10, n_informative=8, random_state=42)
    df = pd.DataFrame(X_raw, columns=[f"feat_{i}" for i in range(10)])
    df["target"] = y_raw
    
    print("[FAULT INJECTED] Severe Data Drift & Missing Feature Values injected!")
    # Ingesting out-of-distribution drift and NaN strings
    df.loc[10:300, "feat_0"] = np.nan  # Unhandled NaNs
    df.loc[301:600, "feat_1"] = df.loc[301:600, "feat_1"] * 99999.0  # Massive scaling drift
    
    # Buggy code fails to impute or scale features
    X = df.drop(columns=["target"])
    y = df["target"]
    
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
    
    print("[Pipeline] Training model on corrupted drifted data...")
    model = RandomForestClassifier(n_estimators=10, random_state=42)
    model.fit(X_train, y_train)
    
    preds = model.predict(X_test)
    acc = accuracy_score(y_test, preds)
    return {"accuracy": float(acc), "loss": float(1.0 - acc), "samples": len(df)}

result = run_ml_pipeline()
""",

    "CODE_RUNTIME_ERROR": """import numpy as np
import pandas as pd
from sklearn.datasets import make_classification

def run_ml_pipeline():
    print("[Pipeline] Ingesting features...")
    X_raw, y_raw = make_classification(n_samples=1000, n_features=5, random_state=42)
    df = pd.DataFrame(X_raw, columns=[f"feat_{i}" for i in range(5)])
    
    print("[FAULT INJECTED] Triggering Runtime Exception in feature aggregation loop...")
    # Unhandled division by zero & missing key access error
    batch_count = 0
    avg_feature = sum(df["feat_0"]) / batch_count  # ZeroDivisionError!
    
    df["target"] = y_raw
    return {"accuracy": 0.0, "loss": 1.0, "samples": len(df)}

result = run_ml_pipeline()
""",

    "NAN_LOSS": """import numpy as np
import pandas as pd
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score

def run_ml_pipeline():
    print("[Pipeline] Training Gradient Boosted Model...")
    X_raw, y_raw = make_classification(n_samples=1000, n_features=5, random_state=42)
    
    print("[FAULT INJECTED] Exploding Gradients resulting in NaN Loss & Inf metrics!")
    loss_weights = np.array([1.0, np.nan, np.inf, 4.0])
    calculated_loss = float(np.mean(loss_weights)) # Returns nan!
    
    if np.isnan(calculated_loss) or np.isinf(calculated_loss):
        raise ValueError(f"CRITICAL MODEL FATAL ERROR: Training Loss evaluated to invalid NaN/Inf ({calculated_loss}). Training aborted.")
        
    return {"accuracy": 0.0, "loss": calculated_loss, "samples": 1000}

result = run_ml_pipeline()
""",

    "OOM_SPIKE": """import numpy as np
import pandas as pd

def run_ml_pipeline():
    print("[Pipeline] Allocating batch buffer for deep learning embeddings...")
    print("[FAULT INJECTED] Memory Spike / Out Of Memory threshold breached!")
    
    # Simulating massive buffer allocation that breaches memory limits
    dummy_huge_array = np.ones((50000, 50000), dtype=np.float64) # ~20GB request simulated
    return {"accuracy": 0.5, "loss": 0.5, "samples": 50000}

result = run_ml_pipeline()
""",

    "MODEL_ACCURACY_DROP": """import numpy as np
import pandas as pd
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score

def run_ml_pipeline():
    print("[Pipeline] Running feature selection and model training...")
    X_raw, y_raw = make_classification(n_samples=1000, n_features=10, n_informative=8, random_state=42)
    
    print("[FAULT INJECTED] Misconfigured hyper-parameters & dropped informative features!")
    # Incorrectly dropping informative features and setting max_depth=1
    X = pd.DataFrame(X_raw).iloc[:, 8:10] # Only keeping 2 weak noise features
    y = y_raw
    
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
    model = RandomForestClassifier(n_estimators=1, max_depth=1, random_state=42)
    model.fit(X_train, y_train)
    
    preds = model.predict(X_test)
    acc = accuracy_score(y_test, preds)
    print(f"[Pipeline Result] Severely Degraded Accuracy: {acc:.4f}")
    return {"accuracy": float(acc), "loss": float(1.0 - acc), "samples": len(X)}

result = run_ml_pipeline()
"""
}

class MLPipelineEngine:
    def __init__(self):
        self.current_code = DEFAULT_PIPELINE_CODE
        self.execution_history = []

    def load_fault_scenario(self, fault_name: str) -> str:
        """
        Loads a pre-defined fault scenario into active pipeline code.
        """
        if fault_name in FAULT_TEMPLATES:
            self.current_code = FAULT_TEMPLATES[fault_name]
        return self.current_code

    def set_custom_code(self, code: str):
        self.current_code = code

    def execute_pipeline(self, code: str = None) -> dict:
        """
        Executes the Python pipeline script in a safe sandboxed environment.
        Captures logs, exceptions, execution time, and memory usage.
        """
        script_to_run = code if code is not None else self.current_code
        self.current_code = script_to_run

        log_capture = io.StringIO()
        old_stdout = sys.stdout
        old_stderr = sys.stderr

        start_time = time.time()
        start_mem = psutil.Process().memory_info().rss / (1024 * 1024)

        status = "HEALTHY"
        error_logs = ""
        result_dict = {"accuracy": 0.0, "loss": 1.0, "samples": 0}

        try:
            sys.stdout = log_capture
            sys.stderr = log_capture
            
            # Local namespace for execution
            exec_globals = {}
            exec(script_to_run, exec_globals)
            
            if "result" in exec_globals and isinstance(exec_globals["result"], dict):
                result_dict = exec_globals["result"]
                acc = result_dict.get("accuracy", 0.0)
                if acc < 0.70:
                    status = "DEGRADED"

        except Exception as e:
            status = "CRITICAL_FAILURE"
            error_logs = traceback.format_exc()
            print("\n=== EXECUTION EXCEPTION TRACEBACK ===", file=log_capture)
            print(error_logs, file=log_capture)

        finally:
            sys.stdout = old_stdout
            sys.stderr = old_stderr

        end_time = time.time()
        end_mem = psutil.Process().memory_info().rss / (1024 * 1024)
        captured_output = log_capture.getvalue()

        # Telemetry metrics
        execution_sec = round(end_time - start_time, 3)
        mem_used_mb = round(max(end_mem, start_mem + np.random.uniform(10, 45)), 1)
        
        telemetry = {
            "status": status,
            "accuracy": float(result_dict.get("accuracy", 0.0)),
            "loss": float(result_dict.get("loss", 1.0)),
            "memory_mb": mem_used_mb,
            "execution_time_sec": execution_sec,
            "samples_processed": result_dict.get("samples", 0),
            "step": len(self.execution_history) + 1
        }

        execution_record = {
            "telemetry": telemetry,
            "logs": captured_output,
            "code": script_to_run,
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
        }
        
        self.execution_history.append(telemetry)
        return execution_record