ai_v/models.py

from extensions import db
from datetime import datetime, timedelta, timezone
from werkzeug.security import generate_password_hash, check_password_hash

# 定义时区常量
UTC_TZ = timezone.utc
BEIJING_TZ = timezone(timedelta(hours=8))

def get_bj_now():
    """获取当前北京时间（无时区信息的 naive datetime）"""
    return datetime.now(BEIJING_TZ).replace(tzinfo=None)

def to_bj_time(dt):
    """返回北京时间
    
    注意：为统一时间处理，现在所有时间存储都直接使用北京时间。
    此函数主要用于兼容性，直接返回输入时间。
    """
    if not dt:
        return None
    # 直接返回（数据库存储的已经是北京时间）
    return dt if dt.tzinfo is None else dt.replace(tzinfo=None)


# 角色与权限的多对多关联表
role_permissions = db.Table('role_permissions',
    db.Column('role_id', db.Integer, db.ForeignKey('roles.id'), primary_key=True),
    db.Column('permission_id', db.Integer, db.ForeignKey('permissions.id'), primary_key=True)
)

class Permission(db.Model):
    __tablename__ = 'permissions'
    id = db.Column(db.Integer, primary_key=True)
    name = db.Column(db.String(50), unique=True, nullable=False)  # 如: 'view_logs', 'manage_users'
    description = db.Column(db.String(100))

class Role(db.Model):
    __tablename__ = 'roles'
    id = db.Column(db.Integer, primary_key=True)
    name = db.Column(db.String(50), unique=True, nullable=False)  # 如: '超级管理员', '普通用户'
    description = db.Column(db.String(100))
    # 角色拥有的权限
    permissions = db.relationship('Permission', secondary=role_permissions, 
                                  backref=db.backref('roles', lazy='dynamic'),
                                  order_by='Permission.id')

class User(db.Model):
    __tablename__ = 'users'
    
    id = db.Column(db.Integer, primary_key=True)
    phone = db.Column(db.String(20), unique=True, nullable=False)
    password_hash = db.Column(db.String(255), nullable=False)
    api_key = db.Column(db.String(255))  # 存储用户的 API Key
    points = db.Column(db.Integer, default=2)  # 账户积分，默认赠送2次试用
    has_used_points = db.Column(db.Boolean, default=False) # 是否使用过积分
    is_banned = db.Column(db.Boolean, default=False) # 账号是否被封禁
    # 关联角色 ID
    role_id = db.Column(db.Integer, db.ForeignKey('roles.id'))
    created_at = db.Column(db.DateTime, default=get_bj_now)

    @property
    def created_at_bj(self):
        return to_bj_time(self.created_at)
    
    # 关系映射
    role = db.relationship('Role', backref=db.backref('users', lazy='dynamic'))

    def has_permission(self, perm_name):
        """动态检查用户是否拥有某项权限"""
        if not self.role:
            return False
        
        # 获取用户拥有的所有权限名称
        perms = [p.name for p in self.role.permissions]
        
        # 核心修复：如果是超级管理员（拥有 manage_system），则豁免所有具体权限检查
        if 'manage_system' in perms:
            return True
            
        return perm_name in perms

    def set_password(self, password):
        self.password_hash = generate_password_hash(password)
        
    def check_password(self, password):
        return check_password_hash(self.password_hash, password)

class SystemDict(db.Model):
    """通用字典管理系统"""
    __tablename__ = 'system_dicts'
    
    id = db.Column(db.Integer, primary_key=True)
    dict_type = db.Column(db.String(50), nullable=False)  # 如: 'ai_model', 'aspect_ratio', 'prompt_tpl'
    label = db.Column(db.String(100), nullable=False)    # 显示名称
    value = db.Column(db.Text, nullable=False)           # 存储值或提示词内容
    cost = db.Column(db.Integer, default=0)              # 消耗积分 (仅针对 ai_model 有效)
    is_active = db.Column(db.Boolean, default=True)
    sort_order = db.Column(db.Integer, default=0)        # 排序权重

class GenerationRecord(db.Model):
    """AI 生成记录"""
    __tablename__ = 'generation_records'
    
    id = db.Column(db.Integer, primary_key=True)
    user_id = db.Column(db.Integer, db.ForeignKey('users.id'), nullable=False)
    prompt = db.Column(db.Text)
    model = db.Column(db.String(100))
    cost = db.Column(db.Integer, default=0) # 消耗积分
    # 存储生成的图片 URL 列表 (JSON 字符串)
    image_urls = db.Column(db.Text)
    created_at = db.Column(db.DateTime, default=get_bj_now)

    @property
    def created_at_bj(self):
        return to_bj_time(self.created_at)
    
    user = db.relationship('User', backref=db.backref('records', lazy='dynamic', order_by='GenerationRecord.created_at.desc()'))

# 用户已读通知关联表
notification_reads = db.Table('notification_reads',
    db.Column('user_id', db.Integer, db.ForeignKey('users.id'), primary_key=True),
    db.Column('notification_id', db.Integer, db.ForeignKey('system_notifications.id'), primary_key=True)
)

class SystemNotification(db.Model):
    """系统全局通知"""
    __tablename__ = 'system_notifications'
    
    id = db.Column(db.Integer, primary_key=True)
    title = db.Column(db.String(200), nullable=False)
    content = db.Column(db.Text, nullable=False)
    is_active = db.Column(db.Boolean, default=True)
    created_at = db.Column(db.DateTime, default=get_bj_now)

    @property
    def created_at_bj(self):
        return to_bj_time(self.created_at)
    
    # 哪些用户已读
    read_by_users = db.relationship('User', secondary=notification_reads, backref=db.backref('read_notifications', lazy='dynamic'))

class Order(db.Model):
    """订单模型"""
    __tablename__ = 'orders'
    
    id = db.Column(db.Integer, primary_key=True)
    out_trade_no = db.Column(db.String(64), unique=True, nullable=False)  # 系统订单号
    user_id = db.Column(db.Integer, db.ForeignKey('users.id'), nullable=False)
    amount = db.Column(db.Numeric(10, 2), nullable=False)  # 金额
    points = db.Column(db.Integer, nullable=False)  # 购买的积分
    status = db.Column(db.String(20), default='PENDING')  # PENDING, PAID, CANCELLED
    trade_no = db.Column(db.String(64))  # 支付宝交易号
    created_at = db.Column(db.DateTime, default=get_bj_now)
    paid_at = db.Column(db.DateTime)

    @property
    def created_at_bj(self):
        return to_bj_time(self.created_at)

    @property
    def paid_at_bj(self):
        return to_bj_time(self.paid_at)
    
    user = db.relationship('User', backref=db.backref('orders', lazy='dynamic', order_by='Order.created_at.desc()'))

class SystemLog(db.Model):
    """系统精细化日志记录"""
    __tablename__ = 'system_logs'
    
    id = db.Column(db.Integer, primary_key=True)
    user_id = db.Column(db.Integer, db.ForeignKey('users.id'), nullable=True)  # 可能没有登录用户
    level = db.Column(db.String(20), nullable=False)    # INFO, WARNING, ERROR, DEBUG
    module = db.Column(db.String(50))                   # 模块名
    message = db.Column(db.Text, nullable=False)        # 日志内容
    extra = db.Column(db.Text)                          # 额外信息的 JSON 字符串
    
    # 请求上下文信息
    ip = db.Column(db.String(50))
    path = db.Column(db.String(255))
    method = db.Column(db.String(10))
    user_agent = db.Column(db.String(255))
    
    created_at = db.Column(db.DateTime, default=get_bj_now)

    @property
    def created_at_bj(self):
        return to_bj_time(self.created_at)
    
    user = db.relationship('User', backref=db.backref('logs', lazy='dynamic', order_by='SystemLog.created_at.desc()'))

class SavedPrompt(db.Model):
    """用户收藏的提示词"""
    __tablename__ = 'saved_prompts'
    
    id = db.Column(db.Integer, primary_key=True)
    user_id = db.Column(db.Integer, db.ForeignKey('users.id'), nullable=False)
    title = db.Column(db.String(100), nullable=False)
    prompt = db.Column(db.Text, nullable=False)
    created_at = db.Column(db.DateTime, default=get_bj_now)

    @property
    def created_at_bj(self):
        return to_bj_time(self.created_at)

    user = db.relationship('User', backref=db.backref('saved_prompts', lazy='dynamic', order_by='SavedPrompt.created_at.desc()'))