from django.db import models
from Device .models import Devices
from Accounts.models import UserProfile

# Create your models here.

class MalwarePrediction(models.Model):
    MODEL_TYPE_CHOICES = [
        (1, 'KNeighborsClassifier'),
        (2, 'RandomForestClassifier'),
        (3, 'XGBClassifier'),
        (4, 'SGDClassifier'),
    ]

    process_name = models.CharField(max_length=255)
    process_class = models.CharField(max_length=50)
    probability_of_malware = models.FloatField()
    predicted_malware = models.CharField(max_length=50)
    model_type = models.IntegerField(choices=MODEL_TYPE_CHOICES, default=1)  # New field to store model type

    def __str__(self):
        return f"{self.process_name} - {self.predicted_malware} - {self.get_model_type_display()}"


class MalwarePredictionsDevice(models.Model):
    device = models.ForeignKey(Devices, on_delete=models.CASCADE,null=True)
    user = models.ForeignKey(UserProfile, on_delete=models.CASCADE , null=True)  # Add this field to reference the user
    file_path = models.FileField(upload_to='malware_predictions/')
    uploaded_at = models.DateTimeField(auto_now_add=True)

    def __str__(self):
        return f"Prediction for {self.device.device_name} by {self.user.user.username} at {self.uploaded_at}"