import gradio as gr
import pickle
import numpy as np

# Modeli yükle
with open("xgboost_optimized_v2.pkl", "rb") as f:  # model/ klasörü kaldırıldı
    model_pkg = pickle.load(f)

# Modeller ve fonksiyonlar
xgb_improved = model_pkg["base_models"]["xgboost_improved"][0]
smart_align = model_pkg["enhancement_functions"]["smart_feature_alignment"]
ensemble_fn = model_pkg["enhancement_functions"]["ensemble_predict"]
unknown_boost_fn = model_pkg["enhancement_functions"]["unknown_category_boost"]
dynamic_adjust_fn = model_pkg["enhancement_functions"]["dynamic_risk_adjustment"]
confidence_fn = model_pkg["enhancement_functions"]["calculate_confidence"]

# Ön işleme bileşenleri
preprocessors = model_pkg["preprocessors"]
scaler = preprocessors["scaler"]
encoders = preprocessors["encoders"]
bert_extractor = preprocessors["bert_extractor"]
label_encoder = preprocessors["label_encoder"]
feature_names = scaler.mean_.shape[0]  # Özellik sayısı

# Tahmin fonksiyonu
def predict(cve_id, cwe_id, description, cvss_score):
    # Dummy girişler (GPU olmadığı için BERT çıkarımı yapılmıyor)
    numeric = np.zeros((1, feature_names))
    bert_feats = np.zeros((1, 768))  # Dummy BERT vektörü
    x_input = np.hstack([numeric, bert_feats])

    # Model tahminleri
    xgb_pred = xgb_improved.predict(x_input)[0]
    ensemble_pred = ensemble_fn(cvss_score, xgb_pred)
    final_pred = unknown_boost_fn(ensemble_pred, unknown_features=2, cvss_score=cvss_score)
    optimized = dynamic_adjust_fn(final_pred, cvss_score, cwe_id, len(description))
    confidence = confidence_fn(2, cvss_score)

    # Kategori sınıflandırması
    category = "Düşük" if optimized <= 3 else "Orta" if optimized <= 7 else "Yüksek"

    return {
        "Tahmini Risk Skoru": round(optimized, 2),
        "Kategori": category,
        "Model Güveni": f"%{round(confidence * 100, 1)}"
    }

# Gradio arayüzü
interface = gr.Interface(
    fn=predict,
    inputs=[
        gr.Text(label="CVE ID"),
        gr.Text(label="CWE ID"),
        gr.Textbox(label="Açıklama"),
        gr.Slider(minimum=0.0, maximum=10.0, step=0.1, label="CVSS Skoru")
    ],
    outputs="json",
    title="Risk Skoru Tahmin Modeli (XGBoost Optimized)",
    description="CVE tanımlarına göre tahmini risk skoru ve kategori"
)

# Uygulamayı başlat
interface.launch()