import joblib
import csv
import sys
from pathlib import Path

import numpy as np
from scipy.sparse import hstack
from sklearn.ensemble import RandomForestClassifier
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
from sklearn.model_selection import train_test_split

sys.path.append(str(Path(__file__).resolve().parent.parent))
from common.features import handcrafted_feature_matrix  # noqa: E402

SQL_KEYWORDS = [
    "select", "union", "drop", "insert", "update", "delete", "or 1=1",
    "sleep(", "xp_cmdshell", "--", "exec ", "'or'", "1=1",
]


def load_data(path="data.csv"):
    texts, labels = [], []
    with open(path, newline="", encoding="utf-8") as f:
        for row in csv.DictReader(f):
            texts.append(row["text"])
            labels.append(int(row["label"]))
    return texts, np.array(labels)


def main():
    texts, labels = load_data()
    X_train_txt, X_test_txt, y_train, y_test = train_test_split(
        texts, labels, test_size=0.25, random_state=42, stratify=labels
    )

    vectorizer = TfidfVectorizer(analyzer="char_wb", ngram_range=(2, 4), max_features=1500)
    Xtr_tfidf = vectorizer.fit_transform(X_train_txt)
    Xte_tfidf = vectorizer.transform(X_test_txt)

    Xtr_hand = handcrafted_feature_matrix(X_train_txt, SQL_KEYWORDS)
    Xte_hand = handcrafted_feature_matrix(X_test_txt, SQL_KEYWORDS)

    X_train = hstack([Xtr_tfidf, Xtr_hand])
    X_test = hstack([Xte_tfidf, Xte_hand])

    clf = RandomForestClassifier(n_estimators=200, random_state=42)
    clf.fit(X_train, y_train)
    y_pred = clf.predict(X_test)

    print("SQL Injection Detection")
    print(f"Accuracy: {accuracy_score(y_test, y_pred):.3f}")
    print(classification_report(y_test, y_pred, target_names=["benign", "malicious"]))
    print("Confusion matrix [[TN FP][FN TP]]:")
    print(confusion_matrix(y_test, y_pred))

    print("\nTest Samples")
    samples = [
        "' OR '1'='1' --",
        "jane.doe@example.com",
        "1; DROP TABLE orders;--",
        "search: blue running shoes size 10",
    ]
    Xd_tfidf = vectorizer.transform(samples)
    Xd_hand = handcrafted_feature_matrix(samples, SQL_KEYWORDS)
    Xd = hstack([Xd_tfidf, Xd_hand])
    preds = clf.predict(Xd)
    for text, label in zip(samples, preds):
        verdict = "MALICIOUS" if label == 1 else "benign"
        print(f"[{verdict}] {text}")


if __name__ == "__main__":
    joblib.dump(vectorizer, "sqli_vectorizer.joblib")
    joblib.dump(clf, "sqli_classifier.joblib")
    print("\nSaved model artifacts. Run: python test_model.py")
    main()