fisk/core/models.py

import stripe
from alpaca.trading.client import TradingClient
from alpaca.trading.enums import OrderSide, TimeInForce
from alpaca.trading.requests import LimitOrderRequest, MarketOrderRequest
from django.conf import settings
from django.contrib.auth.models import AbstractUser
from django.db import models
from serde import ValidationError

from core.lib.customers import get_or_create, update_customer_fields
from core.lib.serde import ccxt_s
from core.util import logs

log = logs.get_logger(__name__)


class Plan(models.Model):
    name = models.CharField(max_length=255, unique=True)
    description = models.CharField(max_length=1024, null=True, blank=True)
    cost = models.IntegerField()
    product_id = models.CharField(max_length=255, unique=True, null=True, blank=True)
    image = models.CharField(max_length=1024, null=True, blank=True)

    def __str__(self):
        return f"{self.name} (£{self.cost})"


class User(AbstractUser):
    # Stripe customer ID
    stripe_id = models.CharField(max_length=255, null=True, blank=True)
    last_payment = models.DateTimeField(null=True, blank=True)
    plans = models.ManyToManyField(Plan, blank=True)
    email = models.EmailField(unique=True)

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._original = self

    def save(self, *args, **kwargs):
        """
        Override the save function to create a Stripe customer.
        """
        if settings.STRIPE_ENABLED:
            if not self.stripe_id:  # stripe ID not stored
                self.stripe_id = get_or_create(
                    self.email, self.first_name, self.last_name
                )

            to_update = {}
            if self.email != self._original.email:
                to_update["email"] = self.email
            if self.first_name != self._original.first_name:
                to_update["first_name"] = self.first_name
            if self.last_name != self._original.last_name:
                to_update["last_name"] = self.last_name

            update_customer_fields(self.stripe_id, **to_update)

        super().save(*args, **kwargs)

    def delete(self, *args, **kwargs):
        if settings.STRIPE_ENABLED:
            if self.stripe_id:
                stripe.Customer.delete(self.stripe_id)
                log.info(f"Deleted Stripe customer {self.stripe_id}")
        super().delete(*args, **kwargs)

    def has_plan(self, plan):
        plan_list = [plan.name for plan in self.plans.all()]
        return plan in plan_list


class Account(models.Model):
    EXCHANGE_CHOICES = (("alpaca", "Alpaca"),)
    user = models.ForeignKey(User, on_delete=models.CASCADE)
    name = models.CharField(max_length=255)
    exchange = models.CharField(choices=EXCHANGE_CHOICES, max_length=255)
    api_key = models.CharField(max_length=255)
    api_secret = models.CharField(max_length=255)
    sandbox = models.BooleanField(default=False)

    def get_account(self):
        trading_client = TradingClient(
            self.api_key, self.api_secret, paper=self.sandbox
        )
        return trading_client.get_account()


class Session(models.Model):
    user = models.ForeignKey(User, on_delete=models.CASCADE)
    request = models.CharField(max_length=255, null=True, blank=True)
    session = models.CharField(max_length=255, null=True, blank=True)
    subscription_id = models.CharField(max_length=255, null=True, blank=True)
    plan = models.ForeignKey(Plan, null=True, blank=True, on_delete=models.CASCADE)


class Hook(models.Model):
    user = models.ForeignKey(User, on_delete=models.CASCADE)
    name = models.CharField(max_length=1024, null=True, blank=True, unique=True)
    hook = models.CharField(max_length=255, unique=True)
    received = models.IntegerField(default=0)


class Trade(models.Model):
    SYMBOL_CHOICES = (
        ("BTC/USD", "Bitcoin/US Dollar"),
        ("LTC/USD", "Litecoin/US Dollar"),
    )
    TYPE_CHOICES = (
        ("market", "Market"),
        ("limit", "Limit"),
    )
    DIRECTION_CHOICES = (
        ("buy", "Buy"),
        ("sell", "Sell"),
    )
    account = models.ForeignKey(Account, on_delete=models.CASCADE)
    hook = models.ForeignKey(Hook, on_delete=models.CASCADE, null=True, blank=True)
    symbol = models.CharField(choices=SYMBOL_CHOICES, max_length=255)
    type = models.CharField(choices=TYPE_CHOICES, max_length=255)
    amount = models.FloatField()
    price = models.FloatField(null=True, blank=True)
    stop_loss = models.FloatField(null=True, blank=True)
    take_profit = models.FloatField(null=True, blank=True)
    status = models.CharField(max_length=255, null=True, blank=True)
    direction = models.CharField(choices=DIRECTION_CHOICES, max_length=255)

    # To populate from the trade
    order_id = models.CharField(max_length=255, null=True, blank=True)
    client_order_id = models.CharField(max_length=255, null=True, blank=True)
    response = models.JSONField(null=True, blank=True)

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._original = self

    def save(self, *args, **kwargs):
        """
        Override the save function to place the trade.
        """
        if self.response is None:
            # the trade is not placed yet
            if self.account.exchange == "alpaca":
                trading_client = TradingClient(
                    self.account.api_key,
                    self.account.api_secret,
                    paper=self.account.sandbox,
                )
                if self.direction == "buy":
                    direction = OrderSide.BUY
                elif self.direction == "sell":
                    direction = OrderSide.SELL
                else:
                    raise Exception("Unknown direction")
                if self.type == "market":

                    market_order_data = MarketOrderRequest(
                        symbol=self.symbol,
                        qty=self.amount,
                        side=OrderSide.BUY,
                        time_in_force=TimeInForce.IOC,
                    )
                    order = trading_client.submit_order(order_data=market_order_data)
                elif self.type == "limit":
                    limit_order_data = LimitOrderRequest(
                        symbol=self.symbol,
                        limit_price=self.price,
                        qty=self.amount,
                        side=direction,
                        time_in_force=TimeInForce.IOC,
                    )
                    order = trading_client.submit_order(order_data=limit_order_data)
                else:
                    raise Exception("Unknown order type")

                print("ORDER", order)

                # self.status = parsed.status
                # self.response = order
        else:
            # there is a trade open
            # get trade
            # update trade
            pass

        super().save(*args, **kwargs)

    def delete(self, *args, **kwargs):
        # close the trade
        super().delete(*args, **kwargs)


class Callback(models.Model):
    hook = models.ForeignKey(Hook, on_delete=models.CASCADE)
    title = models.CharField(max_length=1024, null=True, blank=True)
    message = models.CharField(max_length=1024, null=True, blank=True)
    period = models.CharField(max_length=255, null=True, blank=True)
    timestamp_sent = models.BigIntegerField(null=True, blank=True)
    timestamp_trade = models.BigIntegerField(null=True, blank=True)
    market_exchange = models.CharField(max_length=255, null=True, blank=True)
    market_item = models.CharField(max_length=255, null=True, blank=True)
    market_currency = models.CharField(max_length=255, null=True, blank=True)
    market_contract = models.CharField(max_length=255, null=True, blank=True)


# class Perms(models.Model):
#     class Meta:
#         permissions = (
#             ("bypass_hashing", "Can bypass field hashing"),  #
#             ("bypass_blacklist", "Can bypass the blacklist"),  #
#             ("bypass_encryption", "Can bypass field encryption"),  #
#             ("bypass_obfuscation", "Can bypass field obfuscation"),  #
#             ("bypass_delay", "Can bypass data delay"),  #
#             ("bypass_randomisation", "Can bypass data randomisation"),  #
#             ("post_irc", "Can post to IRC"),
#             ("post_discord", "Can post to Discord"),
#             ("query_search", "Can search with query strings"),  #
#             ("use_insights", "Can use the Insights page"),
#             ("index_int", "Can use the internal index"),
#             ("index_meta", "Can use the meta index"),
#             ("restricted_sources", "Can access restricted sources"),
#         )