ai_v/models.py

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

# 角色与权限的多对多关联表
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'))

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次试用
    # 关联角色 ID
    role_id = db.Column(db.Integer, db.ForeignKey('roles.id'))
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    
    # 关系映射
    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))
    # 存储生成的图片 URL 列表 (JSON 字符串)
    image_urls = db.Column(db.Text)
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    
    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=datetime.utcnow)
    
    # 哪些用户已读
    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=datetime.utcnow)
    paid_at = db.Column(db.DateTime)
    
    user = db.relationship('User', backref=db.backref('orders', lazy='dynamic', order_by='Order.created_at.desc()'))