from django.db import models
from django.contrib.auth.models import User

class UserProfile(models.Model):
    USER_ROLE_CHOICES = (
        ('Free', 'Free'),
        ('Pro', 'Pro'),
        ('Admin', 'Admin'),
    )
    user = models.OneToOneField(User, on_delete=models.CASCADE, related_name='profile')
    role = models.CharField(max_length=10, choices=USER_ROLE_CHOICES, default='Free')

    def __str__(self):
        return f'{self.user.username} - {self.role}'

class Signal(models.Model):
    SIGNAL_CHOICES = (
        ('BUY', 'Buy'),
        ('SELL', 'Sell'),
        ('WAIT', 'Wait'),
    )
    symbol = models.CharField(max_length=20)
    signal = models.CharField(max_length=4, choices=SIGNAL_CHOICES)
    confidence = models.FloatField()
    timestamp = models.DateTimeField(auto_now_add=True)

    def __str__(self):
        return f'{self.timestamp} - {self.symbol} - {self.signal}'

class Trade(models.Model):
    user = models.ForeignKey(User, on_delete=models.CASCADE)
    symbol = models.CharField(max_length=20)
    entry_price = models.FloatField()
    exit_price = models.FloatField(null=True, blank=True)
    entry_timestamp = models.DateTimeField()
    exit_timestamp = models.DateTimeField(null=True, blank=True)
    signal = models.ForeignKey(Signal, on_delete=models.CASCADE)

    def __str__(self):
        return f'{self.user.username} - {self.symbol}'

class Backtest(models.Model):
    user = models.ForeignKey(User, on_delete=models.CASCADE)
    start_date = models.DateTimeField()
    end_date = models.DateTimeField()
    results = models.JSONField()
    timestamp = models.DateTimeField(auto_now_add=True)

    def __str__(self):
        return f'{self.user.username} - {self.timestamp}'