"""
数据库模型和操作
"""
import sqlite3
import json
from datetime import datetime
from contextlib import contextmanager
from config import Config

class Database:
    def __init__(self, db_path=None):
        self.db_path = db_path or Config.DATABASE
        self.init_db()
    
    @contextmanager
    def get_connection(self):
        """获取数据库连接"""
        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row
        try:
            yield conn
        finally:
            conn.close()
    
    def init_db(self):
        """初始化数据库"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            
            # 内容库表
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS articles (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    product_names TEXT NOT NULL,
                    search_title TEXT,
                    category TEXT,
                    keywords TEXT,
                    summary TEXT,
                    content TEXT,
                    source TEXT,
                    url TEXT,
                    fetch_date DATETIME DEFAULT CURRENT_TIMESTAMP,
                    created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                    updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            
            # 为旧表添加字段（如果不存在）
            try:
                cursor.execute('ALTER TABLE articles ADD COLUMN search_title TEXT')
            except:
                pass
            
            # 待处理产品列表
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS pending_products (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    product_name TEXT NOT NULL UNIQUE,
                    category TEXT,
                    subcategory TEXT,
                    priority INTEGER DEFAULT 0,
                    source TEXT,
                    created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                    updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            
            # 处理中产品列表
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS processing_products (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    product_name TEXT NOT NULL UNIQUE,
                    category TEXT,
                    subcategory TEXT,
                    status TEXT DEFAULT 'processing',
                    started_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                    error_message TEXT
                )
            ''')
            
            # 处理历史记录
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS process_history (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    product_name TEXT NOT NULL,
                    category TEXT,
                    subcategory TEXT,
                    status TEXT,
                    review_id TEXT,
                    submitted_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                    details TEXT
                )
            ''')
            
            # 任务配置表
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS task_configs (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    name TEXT NOT NULL UNIQUE,
                    config TEXT,
                    enabled INTEGER DEFAULT 1,
                    created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                    updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            
            # 系统配置表
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS system_config (
                    key TEXT PRIMARY KEY,
                    value TEXT,
                    updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            
            # 失败的URL记录表
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS failed_urls (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    url TEXT NOT NULL,
                    title TEXT,
                    error_message TEXT,
                    retry_count INTEGER DEFAULT 0,
                    status TEXT DEFAULT 'failed',
                    created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                    last_retry_at DATETIME,
                    source TEXT DEFAULT 'search'
                )
            ''')
            
            # 搜索缓存表
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS search_cache (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    keyword TEXT NOT NULL,
                    engine TEXT DEFAULT 'bing_cn',
                    results TEXT NOT NULL,
                    result_count INTEGER,
                    created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                    expires_at DATETIME,
                    UNIQUE(keyword, engine)
                )
            ''')
            
            conn.commit()

# ========== 内容库操作 ==========
    def add_article(self, product_names, category, keywords, summary, content, source, url=None, search_title=None):
        """添加文章到内容库"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                INSERT INTO articles (product_names, search_title, category, keywords, summary, content, source, url)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
            ''', (json.dumps(product_names, ensure_ascii=False), search_title,
                  category, json.dumps(keywords, ensure_ascii=False), summary, content, source, url))
            conn.commit()
            return cursor.lastrowid
    
    def search_articles(self, keyword, category=None):
        """搜索文章"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            if category:
                cursor.execute('''
                    SELECT * FROM articles 
                    WHERE (product_names LIKE ? OR keywords LIKE ? OR summary LIKE ? OR content LIKE ? OR url LIKE ?)
                    AND category = ?
                    ORDER BY fetch_date DESC
                ''', (f'%{keyword}%', f'%{keyword}%', f'%{keyword}%', f'%{keyword}%', f'%{keyword}%', category))
            else:
                cursor.execute('''
                    SELECT * FROM articles 
                    WHERE product_names LIKE ? OR keywords LIKE ? OR summary LIKE ? OR content LIKE ? OR url LIKE ?
                    ORDER BY fetch_date DESC
                ''', (f'%{keyword}%', f'%{keyword}%', f'%{keyword}%', f'%{keyword}%', f'%{keyword}%'))
            return [dict(row) for row in cursor.fetchall()]
    
    def get_article_by_id(self, article_id):
        """获取文章详情"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('SELECT * FROM articles WHERE id = ?', (article_id,))
            row = cursor.fetchone()
            return dict(row) if row else None
    
    def get_all_articles(self, limit=100, offset=0):
        """获取所有文章"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('SELECT * FROM articles ORDER BY fetch_date DESC LIMIT ? OFFSET ?', (limit, offset))
            return [dict(row) for row in cursor.fetchall()]
    
    def delete_article(self, article_id):
        """删除文章"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('DELETE FROM articles WHERE id = ?', (article_id,))
            conn.commit()
            return cursor.rowcount > 0
    
    def get_articles_count(self):
        """获取文章总数"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('SELECT COUNT(*) as count FROM articles')
            row = cursor.fetchone()
            return row['count'] if row else 0

# ========== 待处理产品操作 ==========
    def add_pending_product(self, product_name, category=None, subcategory=None, priority=0, source='manual'):
        """添加待处理产品"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            try:
                cursor.execute('''
                    INSERT INTO pending_products (product_name, category, subcategory, priority, source)
                    VALUES (?, ?, ?, ?, ?)
                ''', (product_name, category, subcategory, priority, source))
                conn.commit()
                return cursor.lastrowid
            except sqlite3.IntegrityError:
                # 产品已存在，更新优先级
                cursor.execute('''
                    UPDATE pending_products 
                    SET priority = MAX(priority, ?), updated_at = CURRENT_TIMESTAMP
                    WHERE product_name = ?
                ''', (priority, product_name))
                conn.commit()
                return None
    
    def get_pending_products(self, limit=10, order_by='priority'):
        """获取待处理产品列表"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            if order_by == 'priority':
                cursor.execute('SELECT * FROM pending_products ORDER BY priority DESC, created_at ASC LIMIT ?', (limit,))
            else:
                cursor.execute('SELECT * FROM pending_products ORDER BY created_at ASC LIMIT ?', (limit,))
            return [dict(row) for row in cursor.fetchall()]
    
    def get_pending_count(self):
        """获取待处理产品数量"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('SELECT COUNT(*) FROM pending_products')
            return cursor.fetchone()[0]
    
    def remove_pending_product(self, product_name):
        """从待处理列表移除产品"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('DELETE FROM pending_products WHERE product_name = ?', (product_name,))
            conn.commit()
            return cursor.rowcount > 0

# ========== 处理中产品操作 ==========
    def start_processing(self, product_name, category, subcategory):
        """开始处理产品"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            try:
                cursor.execute('''
                    INSERT INTO processing_products (product_name, category, subcategory, status)
                    VALUES (?, ?, ?, 'processing')
                ''', (product_name, category, subcategory))
                conn.commit()
                return cursor.lastrowid
            except sqlite3.IntegrityError:
                return None
    
    def finish_processing(self, product_name, status='completed', error_message=None):
        """完成处理"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('DELETE FROM processing_products WHERE product_name = ?', (product_name,))
            conn.commit()
            return cursor.rowcount > 0
    
    def get_processing_products(self):
        """获取处理中的产品"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('SELECT * FROM processing_products')
            return [dict(row) for row in cursor.fetchall()]

# ========== 处理历史操作 ==========
    def add_process_history(self, product_name, category, subcategory, status, review_id=None, details=None):
        """添加处理历史"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                INSERT INTO process_history (product_name, category, subcategory, status, review_id, details)
                VALUES (?, ?, ?, ?, ?, ?)
            ''', (product_name, category, subcategory, status, review_id, 
                  json.dumps(details, ensure_ascii=False) if details else None))
            conn.commit()
            return cursor.lastrowid
    
    def get_process_history(self, limit=100):
        """获取处理历史"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('SELECT * FROM process_history ORDER BY submitted_at DESC LIMIT ?', (limit,))
            return [dict(row) for row in cursor.fetchall()]
    
    def get_history_by_product(self, product_name):
        """获取指定产品的处理历史"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('SELECT * FROM process_history WHERE product_name = ? ORDER BY submitted_at DESC', (product_name,))
            return [dict(row) for row in cursor.fetchall()]

# ========== 任务配置操作 ==========
    def save_task_config(self, name, config):
        """保存任务配置"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                INSERT OR REPLACE INTO task_configs (name, config, updated_at)
                VALUES (?, ?, CURRENT_TIMESTAMP)
            ''', (name, json.dumps(config, ensure_ascii=False)))
            conn.commit()
    
    def get_task_config(self, name):
        """获取任务配置"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('SELECT * FROM task_configs WHERE name = ?', (name,))
            row = cursor.fetchone()
            if row:
                result = dict(row)
                result['config'] = json.loads(result['config'])
                return result
            return None

# ========== 系统配置操作 ==========
    def get_system_config(self, key, default=None):
        """获取系统配置"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('SELECT value FROM system_config WHERE key = ?', (key,))
            row = cursor.fetchone()
            return row['value'] if row else default
    
    def set_system_config(self, key, value):
        """设置系统配置"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                INSERT OR REPLACE INTO system_config (key, value, updated_at)
                VALUES (?, ?, CURRENT_TIMESTAMP)
            ''', (key, value))
            conn.commit()

# ========== 失败URL操作 ==========
    def add_failed_url(self, url, title=None, error_message=None, source='search'):
        """添加失败的URL"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            # 先检查是否已存在
            cursor.execute('SELECT id, retry_count FROM failed_urls WHERE url = ?', (url,))
            existing = cursor.fetchone()
            
            if existing:
                # 更新重试次数和错误信息
                cursor.execute('''
                    UPDATE failed_urls 
                    SET error_message = ?, last_retry_at = CURRENT_TIMESTAMP, retry_count = retry_count + 1
                    WHERE url = ?
                ''', (error_message, url))
            else:
                # 新增失败记录
                cursor.execute('''
                    INSERT INTO failed_urls (url, title, error_message, source)
                    VALUES (?, ?, ?, ?)
                ''', (url, title, error_message, source))
            conn.commit()
            return cursor.lastrowid
    
    def get_failed_urls(self, limit=100, status='failed'):
        """获取失败的URL列表"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                SELECT * FROM failed_urls 
                WHERE status = ?
                ORDER BY created_at DESC 
                LIMIT ?
            ''', (status, limit))
            return [dict(row) for row in cursor.fetchall()]
    
    def get_failed_url_count(self):
        """获取失败URL数量"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('SELECT COUNT(*) FROM failed_urls WHERE status = "failed"')
            return cursor.fetchone()[0]
    
    def mark_url_success(self, url):
        """标记URL为成功（已处理）"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                UPDATE failed_urls 
                SET status = 'success', last_retry_at = CURRENT_TIMESTAMP
                WHERE url = ?
            ''', (url,))
            conn.commit()
            return cursor.rowcount > 0
    
    def delete_failed_url(self, url_id):
        """删除失败URL记录"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('DELETE FROM failed_urls WHERE id = ?', (url_id,))
            conn.commit()
            return cursor.rowcount > 0
    
    def clear_failed_urls(self):
        """清空所有失败URL记录"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('DELETE FROM failed_urls WHERE status = "failed"')
            conn.commit()
            return cursor.rowcount

# ========== 搜索缓存操作 ==========
    def save_search_cache(self, keyword, engine, results, expire_days=7):
        """保存搜索结果缓存"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                INSERT OR REPLACE INTO search_cache (keyword, engine, results, result_count, created_at, expires_at)
                VALUES (?, ?, ?, ?, CURRENT_TIMESTAMP, datetime('now', '+' || ? || ' days'))
            ''', (keyword, engine, json.dumps(results, ensure_ascii=False), len(results), expire_days))
            conn.commit()
            return cursor.lastrowid
    
    def get_search_cache(self, keyword, engine='bing_cn'):
        """获取搜索结果缓存"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                SELECT results, result_count, created_at, expires_at 
                FROM search_cache 
                WHERE keyword = ? AND engine = ? AND expires_at > datetime('now')
            ''', (keyword, engine))
            row = cursor.fetchone()
            if row:
                return {
                    'results': json.loads(row['results']),
                    'count': row['result_count'],
                    'cached_at': row['created_at'],
                    'expires_at': row['expires_at']
                }
            return None
    
    def clear_expired_cache(self):
        """清理过期缓存"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('DELETE FROM search_cache WHERE expires_at <= datetime("now")')
            conn.commit()
            return cursor.rowcount
    
    def clear_all_cache(self):
        """清空所有缓存"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('DELETE FROM search_cache')
            conn.commit()
            return cursor.rowcount

# 全局数据库实例
db = Database()