import json
from pathlib import Path

try:
    from sklearn.feature_extraction.text import CountVectorizer
    from sklearn.naive_bayes import MultinomialNB
    from sklearn.pipeline import make_pipeline
    from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
except ImportError:
    pass  # Allow import to fail if sklearn is not installed, caught at runtime

class InsufficientDataError(ValueError):
    pass

def _load_preprocessed_data(version_dir):
    tsv_path = Path(version_dir) / "preprocessed.tsv"
    if not tsv_path.exists():
        raise FileNotFoundError(f"Missing preprocessed.tsv in {version_dir}")
        
    docs = []
    labels = []
    
    lines = tsv_path.read_text(encoding="utf-8").splitlines()
    if not lines:
        return docs, labels
        
    if lines[0].startswith("label\ttext"):
        lines = lines[1:]
        
    for line in lines:
        parts = line.split("\t")
        if len(parts) >= 2:
            labels.append(parts[0].strip())
            # Join tokens back to strings for sklearn
            docs.append(" ".join(parts[1].split()))
            
    return docs, labels

def _get_version_info(version_dir):
    manifest_path = Path(version_dir) / "manifest.json"
    if manifest_path.exists():
        try:
            m = json.loads(manifest_path.read_text(encoding="utf-8"))
            vid = m.get("version_id", Path(version_dir).name)
            return "V?", vid[:6]
        except Exception:
            pass
    name = Path(version_dir).name
    return "V?", name[:6] if len(name) >= 6 else name

def detect_dataset_shift(old_metrics, new_metrics, thresholds=None):
    if thresholds is None:
        thresholds = {"low": 0.02, "moderate": 0.05, "critical": 0.10}
        
    metrics_to_check = ["accuracy", "precision", "recall", "f1"]
    
    # Add safe fallbacks just in case the JSON keys differ
    def get_val(m_dict, key):
        if key in m_dict: return m_dict[key]
        if key.capitalize() in m_dict: return m_dict[key.capitalize()]
        return None
        
    metrics_to_check = []
    for m in ["accuracy", "precision", "recall", "f1"]:
        if get_val(old_metrics, m) is not None and get_val(new_metrics, m) is not None:
            metrics_to_check.append(m)
            
    if not metrics_to_check:
        raise InsufficientDataError("No common metrics found to compare between old and new versions.")
        
    shift_report = {
        "evaluated_metrics": metrics_to_check,
        "per_metric_drops": {},
        "per_metric_severity": {},
        "overall_severity": "STABLE",
        "recommendation": ""
    }
    
    identical = True
    for m in metrics_to_check:
        if abs(get_val(old_metrics, m) - get_val(new_metrics, m)) > 1e-7:
            identical = False
            break
            
    if identical:
        shift_report["recommendation"] = "Model performs equivalently on new dataset. Retraining not required."
        for m in metrics_to_check:
            shift_report["per_metric_drops"][m] = 0.0
            shift_report["per_metric_severity"][m] = "STABLE"
        return shift_report
        
    old_classes = set(old_metrics.get("classes", []))
    new_classes = set(new_metrics.get("classes", []))
    unseen_classes = new_classes - old_classes if old_classes else set()
    
    if unseen_classes:
        shift_report["unseen_classes"] = list(unseen_classes)
        shift_report["overall_severity"] = "CRITICAL"
        shift_report["recommendation"] = "Significant dataset shift detected. Model cannot explain new data variance. Retraining recommended."
        for m in metrics_to_check:
            drop = get_val(old_metrics, m) - get_val(new_metrics, m)
            shift_report["per_metric_drops"][m] = drop
            shift_report["per_metric_severity"][m] = "CRITICAL"
        return shift_report
        
    worst_severity_level = 0
    severity_map = {0: "STABLE", 1: "MONITOR", 2: "WARNING", 3: "CRITICAL"}
    
    for m in metrics_to_check:
        old_val = get_val(old_metrics, m)
        new_val = get_val(new_metrics, m)
        drop = old_val - new_val
        
        shift_report["per_metric_drops"][m] = max(0.0, drop)
        
        if drop >= thresholds["critical"]:
            shift_report["per_metric_severity"][m] = "CRITICAL"
            worst_severity_level = max(worst_severity_level, 3)
        elif drop >= thresholds["moderate"]:
            shift_report["per_metric_severity"][m] = "WARNING"
            worst_severity_level = max(worst_severity_level, 2)
        elif drop >= thresholds["low"]:
            shift_report["per_metric_severity"][m] = "MONITOR"
            worst_severity_level = max(worst_severity_level, 1)
        else:
            shift_report["per_metric_severity"][m] = "STABLE"
            
    shift_report["overall_severity"] = severity_map[worst_severity_level]
    
    if worst_severity_level == 3:
        shift_report["recommendation"] = "Significant dataset shift detected. Model cannot explain new data variance. Retraining recommended."
    elif worst_severity_level == 2:
        shift_report["recommendation"] = "Moderate performance drop. Consider retraining if next version confirms trend."
    elif worst_severity_level == 1:
        shift_report["recommendation"] = "Minor performance drop detected. Monitor next version before deciding."
    else:
        shift_report["recommendation"] = "Model performs equivalently on new dataset. Retraining not required."
        
    return shift_report

def _parse_text_metrics(text):
    import re
    metrics = {}
    patterns = {
        "accuracy": r"accuracy.*?([\d]*\.[\d]+|[\d]+%)",
        "precision": r"precision.*?([\d]*\.[\d]+|[\d]+%)",
        "recall": r"recall.*?([\d]*\.[\d]+|[\d]+%)",
        "f1": r"f1.*?([\d]*\.[\d]+|[\d]+%)"
    }
    
    for key, pattern in patterns.items():
        match = re.search(pattern, text, re.IGNORECASE)
        if match:
            try:
                val_str = match.group(1).replace('%', '')
                val = float(val_str)
                if '%' in match.group(1) or val > 1.0:
                    val = val / 100.0
                metrics[key] = val
            except ValueError:
                pass
    return metrics

def _get_baseline_metrics(old_dir):
    # Search for explicitly uploaded results file (json, txt, csv, etc.)
    # We use a broad glob to catch "results.txt", "training_results.json", etc.
    for results_path in old_dir.glob("*result*"):
        if results_path.exists():
            text = results_path.read_text(encoding="utf-8", errors="ignore")
            
            if results_path.suffix == ".json":
                try:
                    return json.loads(text)
                except Exception:
                    pass
            
            # If not json, or json parsing failed, attempt Regex text mining
            parsed_metrics = _parse_text_metrics(text)
            if parsed_metrics:
                return parsed_metrics
                
    # Fallback to internal tracker
    eval_path = old_dir / "model_evaluation.json"
    if eval_path.exists():
         return json.loads(eval_path.read_text(encoding="utf-8"))
         
    raise FileNotFoundError(
         f"No baseline metrics found in {old_dir.name}. "
         "Make sure to attach results via 'marco attach-results' or 'marco upload -r'."
    )

def run_full_drift_check(old_version_dir, new_version_dir):
    old_dir = Path(old_version_dir)
    new_dir = Path(new_version_dir)
    
    # 1. Take metrics from the results file and use as baseline
    old_metrics = _get_baseline_metrics(old_dir)
    
    # 2. Train naive bayes model on old data on-the-fly (no saving/joblib needed!)
    old_docs, old_labels = _load_preprocessed_data(old_dir)
    if not old_docs:
        raise InsufficientDataError("Old version has no data.")
        
    model = make_pipeline(CountVectorizer(), MultinomialNB(alpha=1.0))
    model.fit(old_docs, old_labels)
    
    # 3. Use new dataset as test data
    new_docs, new_labels = _load_preprocessed_data(new_dir)
    if not new_docs:
        raise InsufficientDataError("New dataset version is empty.")
        
    predictions = model.predict(new_docs)
    
    # 4. Find new model metrics
    new_metrics = {
        "accuracy": float(accuracy_score(new_labels, predictions)),
        "precision": float(precision_score(new_labels, predictions, average='macro', zero_division=0)),
        "recall": float(recall_score(new_labels, predictions, average='macro', zero_division=0)),
        "f1": float(f1_score(new_labels, predictions, average='macro', zero_division=0)),
        "classes": sorted(list(set(new_labels) | set(predictions)))
    }
    
    # 5. Compare side-by-side
    shift_report = detect_dataset_shift(old_metrics, new_metrics)
    
    # Print formatted output
    v1_label, v1_hash = _get_version_info(old_dir)
    v2_label, v2_hash = _get_version_info(new_dir)
    
    print("MODEL DRIFT ANALYSIS")
    print("════════════════════════════════════════")
    print(f"Trained on:   {v1_label}  ({v1_hash})")
    print(f"Evaluated on: {v2_label}  ({v2_hash})")
    print("════════════════════════════════════════")
    print(f"{'Metric':<12} {v1_label + ' (baseline)':<14} {v2_label + ' (new)':<10} {'Drop':<7} {'Severity':<10}")
    print("────────────────────────────────────────────────────────")
    
    metric_keys = [m.capitalize() for m in shift_report["evaluated_metrics"]]
    
    for title in metric_keys:
        key = title.lower()
        old_val = old_metrics.get(key, old_metrics.get(title, 0.0))
        new_val = new_metrics.get(key, 0.0)
        drop = old_val - new_val
        
        severity = shift_report["per_metric_severity"].get(key, "STABLE")
        if shift_report.get("unseen_classes"):
            severity = "CRITICAL"
        elif shift_report["overall_severity"] == "STABLE":
            severity = "STABLE"
            
        pct_str = f"{-drop * 100:.1f}%"
        if drop <= 0 and drop != 0.0:
             pct_str = f"+{-drop * 100:.1f}%"
             
        print(f"{title:<12} {old_val:<14.3f} {new_val:<10.3f} {pct_str:<7} {severity:<10}")

    print("════════════════════════════════════════")
    overall = shift_report["overall_severity"]
    if shift_report.get("unseen_classes") or overall == "CRITICAL":
        overall = "RETRAIN"
        
    print(f"Overall: {overall}")
    
    import textwrap
    wrapped = textwrap.fill(f'"{shift_report["recommendation"]}"', width=56)
    print(wrapped)
