import numpy as np
import pandas as pd
import logging
import subprocess
import os
import sys
import requests

# Add the project root to the path so we can import from src
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))

from src.models.train import load_real_data

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

DRIFT_THRESHOLD = 0.05  # Significance level (alpha = 0.05) for Kolmogorov-Smirnov test

def send_alert(message):
    """
    Sends a notification alert if a WEBHOOK_URL is provided in the environment.
    """
    webhook_url = os.environ.get("ALERT_WEBHOOK_URL")
    if webhook_url:
        try:
            payload = {"text": f"📢 *MLOps Alert*: {message}"}
            requests.post(webhook_url, json=payload, timeout=5)
            logging.info("Alert notification sent successfully.")
        except Exception as e:
            logging.error(f"Failed to send alert: {e}")
    else:
        logging.info("No ALERT_WEBHOOK_URL found. Skipping notification.")

def fetch_latest_data(simulate_shock=False):
    """
    Simulates fetching the latest 7 days of data from the scraper.
    """
    logging.info("Fetching latest market data...")
    dates = pd.date_range(start='2026-01-01', periods=7, freq='D')
    
    # Baseline latest price would normally be around 250 (from the end of 2025 trend)
    base_price = 250
    
    if simulate_shock:
        logging.warning("Simulating an economic shock (+30% inflation) in the incoming data!")
        # Artificially inflate by 30%
        base_price = base_price * 1.30
        
    noise = np.random.normal(0, 5, 7)
    prices = base_price + noise
    
    df = pd.DataFrame({'Date': dates, 'Price': prices})
    df.set_index('Date', inplace=True)
    return df

def calculate_drift(new_data, baseline_df):
    """
    Calculates statistical concept drift using the Kolmogorov-Smirnov (KS) test.
    The KS test checks whether the distribution of recent prices has drifted
    significantly compared to the baseline training distribution.
    """
    from scipy.stats import ks_2samp
    
    new_prices = new_data['Price'].values
    baseline_prices = baseline_df['Price'].values
    
    # Run Kolmogorov-Smirnov test comparing the two price distributions
    ks_stat, p_value = ks_2samp(baseline_prices, new_prices)
    
    latest_mean = new_prices.mean()
    drift_detected = bool(p_value < DRIFT_THRESHOLD)
    
    return float(ks_stat), float(p_value), drift_detected, latest_mean

def main(simulate_shock=False):
    # 1. Get Baseline Data
    logging.info("Loading baseline training data to calculate baseline mean...")
    baseline_df = load_real_data()
    # Exclude test set (last 20%) to match training data
    train_size = int(len(baseline_df) * 0.8)
    train_df = baseline_df.iloc[:train_size]
    baseline_mean = train_df['Price'].mean()
    logging.info(f"Baseline Mean Price (Training Data): {baseline_mean:.2f} LKR")
    
    # 2. Fetch Latest Data
    latest_data = fetch_latest_data(simulate_shock=simulate_shock)
    
    # 3. Calculate Drift
    ks_stat, p_value, drift_detected, latest_mean = calculate_drift(latest_data, train_df)
    logging.info(f"Latest 7-Day Mean Price: {latest_mean:.2f} LKR")
    logging.info(f"KS Statistic: {ks_stat:.4f}, p-value: {p_value:.4f}")
    logging.info(f"Concept Drift Detected: {drift_detected}")
    
    # 4. Save drift status to JSON for Streamlit dashboard
    import json
    from datetime import datetime
    
    retrain_triggered = drift_detected
    
    status_data = {
        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        "baseline_mean": round(float(baseline_mean), 2),
        "latest_mean": round(float(latest_mean), 2),
        "drift_score": round(float(ks_stat), 4),
        "p_value": round(float(p_value), 4),
        "drift_threshold": round(float(DRIFT_THRESHOLD), 4),
        "drift_detected": drift_detected,
        "retrain_triggered": retrain_triggered,
        "simulation_mode": simulate_shock
    }
    
    os.makedirs("data/processed", exist_ok=True)
    try:
        with open("data/processed/drift_status.json", "w") as f:
            json.dump(status_data, f, indent=4)
        logging.info("Drift status successfully written to data/processed/drift_status.json")
    except Exception as e:
        logging.error(f"Failed to write drift status: {e}")
        
    # 5. Check against threshold
    if drift_detected:
        alert_msg = f"🚨 Concept Drift Detected! p-value ({p_value:.4f}) is below significance threshold ({DRIFT_THRESHOLD})."
        logging.error(alert_msg)
        send_alert(alert_msg)
        
        logging.info("Initiating automatic model retraining pipeline...")
        
        # Trigger retraining
        train_script_path = os.path.join(os.path.dirname(__file__), 'train.py')
        # Setting CWD for subprocess to project root so MLFlow logs to the correct mlruns folder
        project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '../..'))
        try:
            subprocess.run([sys.executable, train_script_path], cwd=project_root, check=True)
            logging.info("✅ Model retraining completed successfully. New model logged to MLflow.")
            send_alert("✅ Model retraining completed successfully. New model is live on DagsHub.")
        except subprocess.CalledProcessError as e:
            err_msg = f"❌ Model retraining failed: {e}"
            logging.error(err_msg)
            send_alert(err_msg)
    else:
        logging.info(f"✅ Data is stable. p-value ({p_value:.4f}) is above significance threshold ({DRIFT_THRESHOLD}). No retraining required.")

if __name__ == "__main__":
    # Check if we should simulate a shock via command line arg
    simulate_shock = "--simulate-shock" in sys.argv
    main(simulate_shock=simulate_shock)