from peewee import CharField, IntegerField, FloatField, BooleanField

from core.database import BaseModel, db

# 策略类型常量
STRATEGY_TYPE_UNCLASSIFIED = 0  # 未分类持仓
STRATEGY_TYPE_GRID = 1          # 网格策略


# 定义Target类，对应targets表
class SFGridTradeTarget(BaseModel):
    stock_code = CharField(unique=True)
    stock_name = CharField()
    current_position = IntegerField()
    grid_index = IntegerField(default=0)
    init_price = FloatField(null=True) # 建仓成本
    grid_match_count = IntegerField(default=0)
    grid_total_profit = FloatField(default=0.0)
    status = IntegerField(default=0)  # -1表示新标的，未完成交易配置，0表示新标的，已完成交易配置，1表示已建初始仓，正常交易中
    enabled = BooleanField(default=False)  # 是否启动交易线程
    strategy_type = IntegerField(default=0)  # 0=未分类, 1=网格策略

    grid_start_price = FloatField(default=10.0) # 基线价格
    grid_size = FloatField(default=1.0)         # 网格价位差
    grid_volume = IntegerField(default=200)     # 网格交易量
    grid_upper_count = IntegerField(default=1)  # 基线价格上方网格数
    grid_lower_count = IntegerField(default=10)  # 基线价格下方网格数

    def targetName(self):
        return f'{self.stock_code}-{self.stock_name}'
    
    def getPriceGrid(self) -> list:
        self.priceGrid: list = []
        # 网格大小，数量
        if self.priceGrid is None or len(self.priceGrid) == 0:
            for i in range(self.grid_upper_count):  # type: ignore 
                upperPrice = self.grid_start_price + (self.grid_upper_count - i) * self.grid_size
                self.priceGrid.append(round(upperPrice, 3))
            
            self.priceGrid.append(self.grid_start_price)
            
            for i in range(self.grid_lower_count): # type: ignore 5
                lowerPrice = self.grid_start_price - (i + 1) * self.grid_size
                self.priceGrid.append(round(lowerPrice, 3))
        
        return self.priceGrid


db.create_tables([SFGridTradeTarget])

# 数据库迁移: 为已有表添加 strategy_type 字段（如果不存在）
try:
    from playhouse.migrate import migrate, SqliteMigrator
    migrator = SqliteMigrator(db)
    migrate(
        migrator.add_column('sfgridtradetarget', 'strategy_type', SFGridTradeTarget.strategy_type),
    )
except Exception:
    # 字段已存在或迁移失败 — 静默跳过
    pass