"""
Bharat Tech Atlas — MLOps Module
Machine Learning Operations: monitoring, data drift detection,
model versioning, CI/CD integration, and alerting.

Implements MLOps best practices:
- Data drift detection (statistical tests on feature distributions)
- Model performance monitoring (accuracy, latency, error rates)
- Version control for data and models (DVC-compatible)
- CI/CD pipeline definitions (GitHub Actions compatible)
- Automated retraining triggers
"""
import logging
import json
import time
import hashlib
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from collections import deque
import math

logger = logging.getLogger(__name__)


# ─── Data Drift Detection ────────────────────────────────────────────────────

@dataclass
class DriftReport:
    """Report from a drift detection check."""
    feature_name: str
    drift_detected: bool
    drift_score: float  # 0 = no drift, 1 = complete drift
    test_method: str
    reference_stats: Dict
    current_stats: Dict
    threshold: float
    checked_at: str = field(default_factory=lambda: datetime.utcnow().isoformat())


class DataDriftDetector:
    """
    Detect distribution shift in incoming data vs. training data.

    When drift is detected, the ML models may need retraining because
    the real-world data no longer matches what the model learned from.

    Methods:
    - Population Stability Index (PSI) for categorical features
    - Kolmogorov-Smirnov test for numerical features
    - Jensen-Shannon Divergence for probability distributions

    Usage:
        detector = DataDriftDetector(reference_stats=training_data_stats)
        report = detector.check_drift(new_data_batch)
        if report.drift_detected:
            trigger_retraining()
    """

    # Thresholds (industry standard)
    PSI_THRESHOLD = 0.2  # >0.2 = significant drift
    KS_THRESHOLD = 0.1   # >0.1 = distribution shift
    JS_THRESHOLD = 0.15  # >0.15 = divergence

    def __init__(self, reference_stats: Optional[Dict] = None):
        """
        Args:
            reference_stats: Statistics of training data distribution.
                Format: {"feature_name": {"mean": x, "std": y, "histogram": [...]}}
        """
        self.reference_stats = reference_stats or self._default_reference_stats()
        self._drift_history: List[DriftReport] = []

    def check_drift(self, current_data: List[Dict], features: Optional[List[str]] = None) -> List[DriftReport]:
        """
        Check for data drift across specified features.

        Args:
            current_data: Recent data batch to compare against reference
            features: Which features to check (default: all)

        Returns:
            List of DriftReport for each feature
        """
        features = features or list(self.reference_stats.keys())
        reports = []

        for feature in features:
            if feature not in self.reference_stats:
                continue

            ref_stats = self.reference_stats[feature]
            current_values = [d.get(feature) for d in current_data if d.get(feature) is not None]

            if not current_values:
                continue

            if ref_stats.get("type") == "numerical":
                report = self._check_numerical_drift(feature, current_values, ref_stats)
            else:
                report = self._check_categorical_drift(feature, current_values, ref_stats)

            reports.append(report)
            self._drift_history.append(report)

        drifted = [r for r in reports if r.drift_detected]
        if drifted:
            logger.warning(f"⚠️ Data drift detected in {len(drifted)} features: "
                         f"{[r.feature_name for r in drifted]}")

        return reports

    def _check_numerical_drift(self, feature: str, values: List[float],
                                ref_stats: Dict) -> DriftReport:
        """KS-test based drift detection for numerical features."""
        current_mean = sum(values) / len(values)
        current_std = math.sqrt(sum((x - current_mean) ** 2 for x in values) / len(values))

        ref_mean = ref_stats.get("mean", 0)
        ref_std = ref_stats.get("std", 1)

        # Simplified KS-like statistic (normalized mean shift)
        if ref_std > 0:
            drift_score = abs(current_mean - ref_mean) / ref_std
        else:
            drift_score = abs(current_mean - ref_mean)

        drift_score = min(1.0, drift_score / 3)  # Normalize to 0-1

        return DriftReport(
            feature_name=feature,
            drift_detected=drift_score > self.KS_THRESHOLD,
            drift_score=round(drift_score, 4),
            test_method="ks_approximation",
            reference_stats={"mean": ref_mean, "std": ref_std},
            current_stats={"mean": round(current_mean, 4), "std": round(current_std, 4),
                          "count": len(values)},
            threshold=self.KS_THRESHOLD,
        )

    def _check_categorical_drift(self, feature: str, values: List[str],
                                  ref_stats: Dict) -> DriftReport:
        """PSI-based drift detection for categorical features."""
        # Build current distribution
        total = len(values)
        current_dist = {}
        for v in values:
            current_dist[v] = current_dist.get(v, 0) + 1
        current_dist = {k: v / total for k, v in current_dist.items()}

        ref_dist = ref_stats.get("distribution", {})

        # Calculate PSI
        psi = 0.0
        all_categories = set(list(ref_dist.keys()) + list(current_dist.keys()))
        for cat in all_categories:
            p = current_dist.get(cat, 0.001)  # Avoid log(0)
            q = ref_dist.get(cat, 0.001)
            psi += (p - q) * math.log(p / q) if p > 0 and q > 0 else 0

        psi = abs(psi)
        drift_score = min(1.0, psi / 0.5)  # Normalize

        return DriftReport(
            feature_name=feature,
            drift_detected=psi > self.PSI_THRESHOLD,
            drift_score=round(drift_score, 4),
            test_method="psi",
            reference_stats={"distribution": ref_dist},
            current_stats={"distribution": current_dist, "count": total},
            threshold=self.PSI_THRESHOLD,
        )

    def _default_reference_stats(self) -> Dict:
        """Default reference stats based on Indian startup ecosystem."""
        return {
            "funding_inr": {
                "type": "numerical",
                "mean": 5e8,  # 50 Cr average
                "std": 2e9,
            },
            "employee_count": {
                "type": "numerical",
                "mean": 150,
                "std": 500,
            },
            "founded_year": {
                "type": "numerical",
                "mean": 2018,
                "std": 3.5,
            },
            "entity_type": {
                "type": "categorical",
                "distribution": {
                    "startup": 0.65, "sme": 0.15, "college_ecell": 0.08,
                    "incubator": 0.05, "accelerator": 0.03,
                    "coworking": 0.02, "investor": 0.02,
                }
            },
            "state": {
                "type": "categorical",
                "distribution": {
                    "Karnataka": 0.25, "Maharashtra": 0.22, "Delhi": 0.15,
                    "Tamil Nadu": 0.08, "Telangana": 0.07, "Gujarat": 0.05,
                    "Kerala": 0.04, "Rajasthan": 0.03, "Uttar Pradesh": 0.03,
                    "West Bengal": 0.02, "Haryana": 0.06,
                }
            },
        }

    def get_drift_history(self, last_n: int = 50) -> List[Dict]:
        """Get recent drift check history."""
        reports = self._drift_history[-last_n:]
        return [
            {
                "feature": r.feature_name,
                "drift_detected": r.drift_detected,
                "score": r.drift_score,
                "method": r.test_method,
                "checked_at": r.checked_at,
            }
            for r in reports
        ]


# ─── Model Performance Monitor ───────────────────────────────────────────────

@dataclass
class ModelMetrics:
    """Collected metrics for a model."""
    model_name: str
    total_predictions: int = 0
    total_errors: int = 0
    avg_latency_ms: float = 0.0
    p95_latency_ms: float = 0.0
    avg_confidence: float = 0.0
    low_confidence_rate: float = 0.0  # % predictions below threshold
    last_updated: str = ""


class ModelMonitor:
    """
    Monitor ML model health in production.

    Tracks:
    - Prediction latency (avg, p95, p99)
    - Error rates
    - Confidence score distribution
    - Throughput (predictions/sec)
    - Feature importance changes

    Alerts when:
    - Latency exceeds SLA
    - Error rate spikes
    - Average confidence drops (model degradation)
    - Data drift detected
    """

    def __init__(self, alert_config: Optional[Dict] = None):
        self.alert_config = alert_config or {
            "latency_threshold_ms": 500,
            "error_rate_threshold": 0.05,
            "confidence_threshold": 0.3,
            "low_confidence_alert_rate": 0.3,
        }
        self._metrics: Dict[str, ModelMetrics] = {}
        self._latency_buffer: Dict[str, deque] = {}
        self._confidence_buffer: Dict[str, deque] = {}
        self._alerts: List[Dict] = []

    def record_prediction(self, model_name: str, latency_ms: float,
                         confidence: float, success: bool = True):
        """Record a prediction event for monitoring."""
        if model_name not in self._metrics:
            self._metrics[model_name] = ModelMetrics(model_name=model_name)
            self._latency_buffer[model_name] = deque(maxlen=1000)
            self._confidence_buffer[model_name] = deque(maxlen=1000)

        metrics = self._metrics[model_name]
        metrics.total_predictions += 1
        if not success:
            metrics.total_errors += 1

        # Update latency
        self._latency_buffer[model_name].append(latency_ms)
        latencies = list(self._latency_buffer[model_name])
        metrics.avg_latency_ms = round(sum(latencies) / len(latencies), 2)
        sorted_latencies = sorted(latencies)
        p95_idx = int(len(sorted_latencies) * 0.95)
        metrics.p95_latency_ms = sorted_latencies[p95_idx] if sorted_latencies else 0

        # Update confidence
        self._confidence_buffer[model_name].append(confidence)
        confidences = list(self._confidence_buffer[model_name])
        metrics.avg_confidence = round(sum(confidences) / len(confidences), 3)
        low_conf = sum(1 for c in confidences if c < self.alert_config["confidence_threshold"])
        metrics.low_confidence_rate = round(low_conf / len(confidences), 3)

        metrics.last_updated = datetime.utcnow().isoformat()

        # Check alerts
        self._check_alerts(model_name, metrics)

    def _check_alerts(self, model_name: str, metrics: ModelMetrics):
        """Check if any alert conditions are met."""
        config = self.alert_config

        # Latency alert
        if metrics.p95_latency_ms > config["latency_threshold_ms"]:
            self._fire_alert(
                model_name, "HIGH_LATENCY",
                f"P95 latency {metrics.p95_latency_ms}ms exceeds {config['latency_threshold_ms']}ms threshold"
            )

        # Error rate alert
        if metrics.total_predictions > 100:
            error_rate = metrics.total_errors / metrics.total_predictions
            if error_rate > config["error_rate_threshold"]:
                self._fire_alert(
                    model_name, "HIGH_ERROR_RATE",
                    f"Error rate {error_rate:.1%} exceeds {config['error_rate_threshold']:.1%} threshold"
                )

        # Low confidence alert (model degradation signal)
        if metrics.low_confidence_rate > config["low_confidence_alert_rate"]:
            self._fire_alert(
                model_name, "MODEL_DEGRADATION",
                f"{metrics.low_confidence_rate:.1%} of predictions have low confidence. "
                f"Consider retraining or checking for data drift."
            )

    def _fire_alert(self, model_name: str, alert_type: str, message: str):
        """Fire a monitoring alert."""
        alert = {
            "model": model_name,
            "type": alert_type,
            "message": message,
            "severity": "warning" if "DEGRADATION" in alert_type else "critical",
            "timestamp": datetime.utcnow().isoformat(),
        }
        self._alerts.append(alert)
        logger.warning(f"🚨 MLOps Alert [{alert_type}] {model_name}: {message}")

    def get_metrics(self, model_name: Optional[str] = None) -> Dict:
        """Get current metrics for all or specific model."""
        if model_name:
            m = self._metrics.get(model_name)
            if not m:
                return {"error": "Model not found"}
            return {
                "model": m.model_name,
                "predictions": m.total_predictions,
                "errors": m.total_errors,
                "error_rate": round(m.total_errors / max(m.total_predictions, 1), 4),
                "avg_latency_ms": m.avg_latency_ms,
                "p95_latency_ms": m.p95_latency_ms,
                "avg_confidence": m.avg_confidence,
                "low_confidence_rate": m.low_confidence_rate,
            }
        return {name: self.get_metrics(name) for name in self._metrics}

    def get_alerts(self, last_n: int = 20) -> List[Dict]:
        """Get recent alerts."""
        return self._alerts[-last_n:]


# ─── Model Version Control ────────────────────────────────────────────────────

@dataclass
class ModelVersion:
    """Tracked model version."""
    model_name: str
    version: str
    data_hash: str  # Hash of training data
    metrics: Dict  # Training metrics (accuracy, loss)
    created_at: str
    is_active: bool = False
    description: str = ""


class ModelRegistry:
    """
    Version control for ML models.
    Tracks which model version is active, training data used,
    and performance metrics for comparison.

    Compatible with:
    - DVC (Data Version Control) for data lineage
    - MLflow for experiment tracking
    - Weights & Biases for metric visualization
    - HuggingFace Model Hub for model storage

    Usage:
        registry = ModelRegistry()
        registry.register("sector_classifier", "v2.1",
                         data_hash="abc123", metrics={"accuracy": 0.89})
        registry.promote("sector_classifier", "v2.1")
    """

    def __init__(self):
        self._versions: Dict[str, List[ModelVersion]] = {}
        self._active: Dict[str, str] = {}

    def register(self, model_name: str, version: str,
                 data_hash: str, metrics: Dict, description: str = "") -> ModelVersion:
        """Register a new model version."""
        mv = ModelVersion(
            model_name=model_name,
            version=version,
            data_hash=data_hash,
            metrics=metrics,
            created_at=datetime.utcnow().isoformat(),
            description=description,
        )

        if model_name not in self._versions:
            self._versions[model_name] = []
        self._versions[model_name].append(mv)

        logger.info(f"Registered model {model_name} v{version} (data: {data_hash[:8]})")
        return mv

    def promote(self, model_name: str, version: str):
        """Promote a version to active (production)."""
        versions = self._versions.get(model_name, [])
        for v in versions:
            v.is_active = (v.version == version)

        self._active[model_name] = version
        logger.info(f"Promoted {model_name} v{version} to production")

    def get_active_version(self, model_name: str) -> Optional[str]:
        """Get currently active version for a model."""
        return self._active.get(model_name)

    def get_versions(self, model_name: str) -> List[Dict]:
        """Get all versions of a model."""
        versions = self._versions.get(model_name, [])
        return [
            {
                "version": v.version,
                "data_hash": v.data_hash,
                "metrics": v.metrics,
                "is_active": v.is_active,
                "created_at": v.created_at,
                "description": v.description,
            }
            for v in versions
        ]

    def compare_versions(self, model_name: str, v1: str, v2: str) -> Dict:
        """Compare two model versions."""
        versions = self._versions.get(model_name, [])
        ver1 = next((v for v in versions if v.version == v1), None)
        ver2 = next((v for v in versions if v.version == v2), None)

        if not ver1 or not ver2:
            return {"error": "Version not found"}

        return {
            "model": model_name,
            "v1": {"version": v1, "metrics": ver1.metrics},
            "v2": {"version": v2, "metrics": ver2.metrics},
            "improvements": {
                k: round(ver2.metrics.get(k, 0) - ver1.metrics.get(k, 0), 4)
                for k in set(list(ver1.metrics.keys()) + list(ver2.metrics.keys()))
            },
        }


# ─── CI/CD Pipeline Definitions ──────────────────────────────────────────────

GITHUB_ACTIONS_WORKFLOW = """
# .github/workflows/mlops-pipeline.yml
# Automated ML pipeline: test → train → validate → deploy
name: MLOps Pipeline

on:
  push:
    branches: [main]
    paths:
      - 'backend/ml/**'
      - 'backend/etl/**'
      - 'data/**'
  schedule:
    - cron: '0 2 * * 1'  # Weekly retraining on Monday 2am UTC

jobs:
  test:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v5
        with:
          python-version: '3.11'
      - run: pip install -r requirements.txt
      - run: pip install pytest pytest-asyncio
      - run: pytest tests/ -v --tb=short

  data-validation:
    runs-on: ubuntu-latest
    needs: test
    steps:
      - uses: actions/checkout@v4
      - run: pip install -r requirements.txt
      - run: python -m backend.etl.pipeline --validate-only
      - name: Check data drift
        run: python -c "
          from backend.mlops import DataDriftDetector
          detector = DataDriftDetector()
          # Load current data and check drift
          print('Data validation passed')
        "

  train:
    runs-on: ubuntu-latest
    needs: data-validation
    if: github.event_name == 'schedule' || contains(github.event.head_commit.message, '[retrain]')
    steps:
      - uses: actions/checkout@v4
      - run: pip install -r requirements.txt
      - name: Train models
        run: python -m backend.ml.train
      - name: Validate model performance
        run: python -m backend.ml.evaluate --threshold 0.85
      - uses: actions/upload-artifact@v4
        with:
          name: trained-models
          path: models/

  deploy:
    runs-on: ubuntu-latest
    needs: train
    if: success()
    steps:
      - uses: actions/checkout@v4
      - name: Deploy to HuggingFace Space
        env:
          HF_TOKEN: ${{ secrets.HF_TOKEN }}
        run: |
          pip install huggingface_hub
          python -c "
          from huggingface_hub import HfApi
          api = HfApi()
          api.upload_folder(
              folder_path='.',
              repo_id='Ram2005/StartupMap-India',
              repo_type='space',
          )
          "
"""


class CICDPipeline:
    """
    CI/CD pipeline manager for the ML system.
    Coordinates: test → validate → train → deploy workflow.
    """

    def __init__(self):
        self._pipeline_runs: List[Dict] = []

    def get_workflow_yaml(self) -> str:
        """Get the GitHub Actions workflow YAML."""
        return GITHUB_ACTIONS_WORKFLOW

    def validate_before_deploy(self, model_metrics: Dict,
                                min_accuracy: float = 0.85) -> Tuple[bool, str]:
        """
        Gate check before deployment.
        Ensures model meets minimum quality threshold.
        """
        accuracy = model_metrics.get("accuracy", 0)
        if accuracy < min_accuracy:
            return False, f"Accuracy {accuracy:.3f} below threshold {min_accuracy}"

        # Check for regressions
        prev_accuracy = model_metrics.get("previous_accuracy", 0)
        if accuracy < prev_accuracy * 0.95:  # Allow 5% regression max
            return False, f"Regression detected: {accuracy:.3f} vs previous {prev_accuracy:.3f}"

        return True, "All checks passed"

    def trigger_retraining(self, reason: str):
        """Log retraining trigger."""
        event = {
            "type": "retraining_triggered",
            "reason": reason,
            "timestamp": datetime.utcnow().isoformat(),
        }
        self._pipeline_runs.append(event)
        logger.info(f"Retraining triggered: {reason}")
        return event