"""
数据库模型和操作
"""
import sqlite3
import os
from datetime import datetime
import math

# 数据库文件路径
DB_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), 'backend', 'database', 'app.db')

def init_database():
    """初始化数据库，创建表结构"""
    # 确保数据库目录存在
    os.makedirs(os.path.dirname(DB_PATH), exist_ok=True)
    
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    
    # 创建提示词表
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS prompts (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            name VARCHAR(100) NOT NULL,
            content TEXT NOT NULL,
            model_type VARCHAR(50) NOT NULL,
            created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
            updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
        )
    ''')
    
    # 创建生成记录表
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS generation_history (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            prompt_text TEXT NOT NULL,
            model_name VARCHAR(50) NOT NULL,
            image_count INTEGER NOT NULL,
            success_count INTEGER DEFAULT 0,
            created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
            status VARCHAR(20) DEFAULT 'pending'
        )
    ''')
    
    conn.commit()
    conn.close()

def get_db_connection():
    """获取数据库连接"""
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row  # 使结果可以像字典一样访问
    return conn

# 提示词相关操作
def get_all_prompts():
    """获取所有提示词"""
    conn = get_db_connection()
    cursor = conn.cursor()
    
    cursor.execute('''
        SELECT id, name, content, model_type, created_at, updated_at
        FROM prompts
        ORDER BY updated_at DESC
    ''')
    
    prompts = []
    for row in cursor.fetchall():
        prompts.append({
            'id': row['id'],
            'name': row['name'],
            'content': row['content'],
            'model_type': row['model_type'],
            'created_at': row['created_at'],
            'updated_at': row['updated_at']
        })
    
    conn.close()
    return prompts

def get_prompt_by_id(prompt_id):
    """根据ID获取提示词"""
    conn = get_db_connection()
    cursor = conn.cursor()
    
    cursor.execute('''
        SELECT id, name, content, model_type, created_at, updated_at
        FROM prompts
        WHERE id = ?
    ''', (prompt_id,))
    
    row = cursor.fetchone()
    conn.close()
    
    if row:
        return {
            'id': row['id'],
            'name': row['name'],
            'content': row['content'],
            'model_type': row['model_type'],
            'created_at': row['created_at'],
            'updated_at': row['updated_at']
        }
    return None

def save_prompt(name, content, model_type):
    """保存新提示词"""
    conn = get_db_connection()
    cursor = conn.cursor()
    
    cursor.execute('''
        INSERT INTO prompts (name, content, model_type, created_at, updated_at)
        VALUES (?, ?, ?, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
    ''', (name, content, model_type))
    
    prompt_id = cursor.lastrowid
    conn.commit()
    conn.close()
    
    return prompt_id

def delete_prompt_by_id(prompt_id):
    """删除提示词"""
    conn = get_db_connection()
    cursor = conn.cursor()
    
    cursor.execute('DELETE FROM prompts WHERE id = ?', (prompt_id,))
    
    deleted_count = cursor.rowcount
    conn.commit()
    conn.close()
    
    return deleted_count > 0

# 生成历史相关操作
def save_generation_history(prompt_text, model_name, image_count, success_count):
    """保存生成历史记录"""
    conn = get_db_connection()
    cursor = conn.cursor()
    
    status = 'completed' if success_count > 0 else 'failed'
    
    cursor.execute('''
        INSERT INTO generation_history 
        (prompt_text, model_name, image_count, success_count, status, created_at)
        VALUES (?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
    ''', (prompt_text, model_name, image_count, success_count, status))
    
    conn.commit()
    conn.close()

def get_generation_history(page=1, limit=20):
    """获取生成历史记录（分页）"""
    conn = get_db_connection()
    cursor = conn.cursor()
    
    # 计算偏移量
    offset = (page - 1) * limit
    
    # 获取总记录数
    cursor.execute('SELECT COUNT(*) as total FROM generation_history')
    total = cursor.fetchone()['total']
    
    # 获取分页数据
    cursor.execute('''
        SELECT id, prompt_text, model_name, image_count, success_count, status, created_at
        FROM generation_history
        ORDER BY created_at DESC
        LIMIT ? OFFSET ?
    ''', (limit, offset))
    
    records = []
    for row in cursor.fetchall():
        records.append({
            'id': row['id'],
            'prompt_text': row['prompt_text'],
            'model_name': row['model_name'],
            'image_count': row['image_count'],
            'success_count': row['success_count'],
            'status': row['status'],
            'created_at': row['created_at']
        })
    
    conn.close()
    
    # 计算总页数
    pages = math.ceil(total / limit) if total > 0 else 1
    
    return {
        'records': records,
        'total': total,
        'pages': pages
    }

# 数据库维护操作
def clear_old_history(days=30):
    """清理旧的历史记录"""
    conn = get_db_connection()
    cursor = conn.cursor()
    
    cursor.execute('''
        DELETE FROM generation_history 
        WHERE created_at < datetime('now', '-{} days')
    '''.format(days))
    
    deleted_count = cursor.rowcount
    conn.commit()
    conn.close()
    
    return deleted_count