import argparse
import sys
from pathlib import Path

import joblib
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import StratifiedKFold, cross_validate
from sklearn.metrics import classification_report, roc_auc_score

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from app.engine.ml_classifier import extract_features, FEATURE_NAMES


def load_samples(data_dir: Path) -> tuple[np.ndarray, np.ndarray]:
    benign_dir = data_dir / "benign"
    malicious_dir = data_dir / "malicious"

    X_list = []
    y_list = []
    skipped = 0

    for label, directory in [(0, benign_dir), (1, malicious_dir)]:
        if not directory.exists():
            print(f"Directory not found: {directory}")
            continue
        for filepath in directory.iterdir():
            if not filepath.is_file():
                continue
            try:
                file_bytes = filepath.read_bytes()
                features = extract_features(file_bytes)
                if features is not None:
                    X_list.append(features)
                    y_list.append(label)
                else:
                    skipped += 1
            except Exception:
                skipped += 1

    print(f"Loaded {len(X_list)} samples ({skipped} skipped)")
    return np.array(X_list), np.array(y_list)


def train(data_dir: str, output_path: str, model_type: str = "rf") -> None:
    X, y = load_samples(Path(data_dir))

    if len(X) < 20:
        print("Not enough samples for training. Need at least 20.")
        sys.exit(1)

    benign_count = int(np.sum(y == 0))
    malicious_count = int(np.sum(y == 1))
    print(f"Benign: {benign_count}, Malicious: {malicious_count}")

    if model_type == "catboost":
        try:
            from catboost import CatBoostClassifier
            model = CatBoostClassifier(
                iterations=500,
                learning_rate=0.05,
                depth=6,
                l2_leaf_reg=5,
                border_count=128,
                class_weights={0: 1.0, 1: 2.0},
                eval_metric="Precision",
                random_seed=42,
                verbose=0,
            )
        except ImportError:
            print("CatBoost not installed, falling back to RandomForest")
            model_type = "rf"

    if model_type == "rf":
        model = RandomForestClassifier(
            n_estimators=100,
            max_depth=10,
            random_state=42,
        )

    cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
    cv_results = cross_validate(
        model, X, y, cv=cv,
        scoring=["precision", "recall", "f1", "roc_auc"],
        return_train_score=False,
    )

    print("\n5-Fold Cross-Validation Results:")
    print(f"  Precision: {cv_results['test_precision'].mean():.4f} (+/- {cv_results['test_precision'].std():.4f})")
    print(f"  Recall:    {cv_results['test_recall'].mean():.4f} (+/- {cv_results['test_recall'].std():.4f})")
    print(f"  F1:        {cv_results['test_f1'].mean():.4f} (+/- {cv_results['test_f1'].std():.4f})")
    print(f"  AUC-ROC:   {cv_results['test_roc_auc'].mean():.4f} (+/- {cv_results['test_roc_auc'].std():.4f})")

    model.fit(X, y)

    output = Path(output_path)
    output.parent.mkdir(parents=True, exist_ok=True)
    joblib.dump(model, output)
    print(f"\nModel saved to {output}")

    if hasattr(model, "feature_importances_"):
        print("\nFeature Importances:")
        importances = model.feature_importances_
        sorted_idx = np.argsort(importances)[::-1]
        for idx in sorted_idx:
            print(f"  {FEATURE_NAMES[idx]:30s} {importances[idx]:.4f}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--data-dir", required=True)
    parser.add_argument("--output", default="app/models/rf_model.pkl")
    parser.add_argument("--model-type", choices=["rf", "catboost"], default="rf")
    args = parser.parse_args()
    train(args.data_dir, args.output, args.model_type)