"""
会话存储模块 - 提供持久化存储功能

功能:
1. 用户会话历史保存到 SQLite 数据库
2. 活跃任务状态保存
3. 服务重启后能恢复会话
"""

import sqlite3
import json
import os
from datetime import datetime
from pathlib import Path
from typing import Optional, List, Dict, Any

# 默认数据库路径
DEFAULT_DB_PATH = str(Path(__file__).resolve().parent / "data" / "session_storage.db")


class SessionStorage:
    """会话存储管理器"""
    
    def __init__(self, db_path: str = DEFAULT_DB_PATH):
        self.db_path = db_path
        self._ensure_db_dir()
        self._init_tables()
    
    def _ensure_db_dir(self):
        """确保数据库目录存在"""
        db_dir = os.path.dirname(self.db_path)
        if db_dir and not os.path.exists(db_dir):
            os.makedirs(db_dir, exist_ok=True)
    
    def _get_connection(self) -> sqlite3.Connection:
        """获取数据库连接"""
        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row
        return conn
    
    def _init_tables(self):
        """初始化数据库表"""
        conn = self._get_connection()
        cursor = conn.cursor()
        
        # 用户会话历史表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS user_sessions (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id TEXT NOT NULL,
                task_id TEXT NOT NULL,
                job_id TEXT NOT NULL,
                prefix TEXT NOT NULL,
                stage TEXT NOT NULL,
                status TEXT NOT NULL,
                timestamp TEXT NOT NULL,
                final_members INTEGER,
                created_at DATETIME DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        # 用户活跃任务表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS user_active_jobs (
                user_id TEXT PRIMARY KEY,
                job_id TEXT NOT NULL,
                prefix TEXT NOT NULL,
                created_at TEXT NOT NULL,
                stages TEXT,  -- JSON 格式存储
                updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        # 运行中任务表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS running_tasks (
                task_id TEXT PRIMARY KEY,
                job_id TEXT NOT NULL,
                prefix TEXT NOT NULL,
                stage TEXT NOT NULL,
                status TEXT NOT NULL,
                timestamp TEXT NOT NULL,
                user_id TEXT NOT NULL,
                final_members INTEGER,
                updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        # 创建索引
        cursor.execute('''
            CREATE INDEX IF NOT EXISTS idx_user_sessions_user_id 
            ON user_sessions(user_id)
        ''')
        cursor.execute('''
            CREATE INDEX IF NOT EXISTS idx_user_sessions_timestamp 
            ON user_sessions(timestamp)
        ''')
        
        conn.commit()
        conn.close()
    
    # ==================== 用户会话历史 ====================
    
    def get_user_history(self, user_id: str, limit: int = 10) -> List[Dict[str, Any]]:
        """获取用户历史任务"""
        conn = self._get_connection()
        cursor = conn.cursor()
        
        cursor.execute('''
            SELECT task_id, job_id, prefix, stage, status, timestamp, final_members
            FROM user_sessions
            WHERE user_id = ?
            ORDER BY timestamp DESC
            LIMIT ?
        ''', (user_id, limit))
        
        rows = cursor.fetchall()
        conn.close()
        
        return [dict(row) for row in reversed(rows)]  # 反转回时间正序
    
    def add_user_task(self, user_id: str, task_info: Dict[str, Any]):
        """添加任务到用户历史"""
        conn = self._get_connection()
        cursor = conn.cursor()
        
        cursor.execute('''
            INSERT INTO user_sessions 
            (user_id, task_id, job_id, prefix, stage, status, timestamp, final_members)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?)
        ''', (
            user_id,
            task_info.get('task_id', ''),
            task_info.get('job_id', ''),
            task_info.get('prefix', ''),
            task_info.get('stage', ''),
            task_info.get('status', ''),
            task_info.get('timestamp', datetime.now().strftime('%Y-%m-%d %H:%M:%S')),
            task_info.get('final_members')
        ))
        
        conn.commit()
        conn.close()
    
    def update_task_status(self, user_id: str, task_id: str, status: str, 
                          final_members: Optional[int] = None):
        """更新任务状态"""
        conn = self._get_connection()
        cursor = conn.cursor()
        
        if final_members is not None:
            cursor.execute('''
                UPDATE user_sessions 
                SET status = ?, final_members = ?
                WHERE user_id = ? AND task_id = ?
            ''', (status, final_members, user_id, task_id))
        else:
            cursor.execute('''
                UPDATE user_sessions 
                SET status = ?
                WHERE user_id = ? AND task_id = ?
            ''', (status, user_id, task_id))
        
        conn.commit()
        conn.close()
    
    def get_latest_task(self, user_id: str) -> Optional[Dict[str, Any]]:
        """获取用户最近一个任务"""
        conn = self._get_connection()
        cursor = conn.cursor()
        
        cursor.execute('''
            SELECT task_id, job_id, prefix, stage, status, timestamp, final_members
            FROM user_sessions
            WHERE user_id = ?
            ORDER BY timestamp DESC
            LIMIT 1
        ''', (user_id,))
        
        row = cursor.fetchone()
        conn.close()
        
        return dict(row) if row else None
    
    # ==================== 用户活跃任务 ====================
    
    def get_user_active_job(self, user_id: str) -> Optional[Dict[str, Any]]:
        """获取用户当前活跃 job"""
        conn = self._get_connection()
        cursor = conn.cursor()
        
        cursor.execute('''
            SELECT user_id, job_id, prefix, created_at, stages
            FROM user_active_jobs
            WHERE user_id = ?
        ''', (user_id,))
        
        row = cursor.fetchone()
        conn.close()
        
        if row:
            job = dict(row)
            job['stages'] = json.loads(job['stages'] or '{}')
            return job
        return None
    
    def set_user_active_job(self, user_id: str, job_info: Dict[str, Any]):
        """设置用户当前活跃 job"""
        conn = self._get_connection()
        cursor = conn.cursor()
        
        stages_json = json.dumps(job_info.get('stages', {}))
        
        cursor.execute('''
            INSERT OR REPLACE INTO user_active_jobs 
            (user_id, job_id, prefix, created_at, stages, updated_at)
            VALUES (?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
        ''', (
            user_id,
            job_info.get('job_id', ''),
            job_info.get('prefix', ''),
            job_info.get('created_at', datetime.now().strftime('%Y-%m-%d %H:%M:%S')),
            stages_json
        ))
        
        conn.commit()
        conn.close()
    
    def update_job_stage(self, user_id: str, stage: str, status: str):
        """更新 job 的阶段状态"""
        conn = self._get_connection()
        cursor = conn.cursor()
        
        # 先获取当前 stages
        cursor.execute('''
            SELECT stages FROM user_active_jobs WHERE user_id = ?
        ''', (user_id,))
        
        row = cursor.fetchone()
        if row:
            stages = json.loads(row['stages'] or '{}')
            stages[stage] = status
            
            cursor.execute('''
                UPDATE user_active_jobs 
                SET stages = ?, updated_at = CURRENT_TIMESTAMP
                WHERE user_id = ?
            ''', (json.dumps(stages), user_id))
            
            conn.commit()
        
        conn.close()
    
    def clear_user_active_job(self, user_id: str):
        """清除用户活跃 job"""
        conn = self._get_connection()
        cursor = conn.cursor()
        
        cursor.execute('''
            DELETE FROM user_active_jobs WHERE user_id = ?
        ''', (user_id,))
        
        conn.commit()
        conn.close()
    
    # ==================== 运行中任务 ====================
    
    def get_running_tasks(self) -> Dict[str, Dict[str, Any]]:
        """获取所有运行中任务"""
        conn = self._get_connection()
        cursor = conn.cursor()
        
        cursor.execute('''
            SELECT task_id, job_id, prefix, stage, status, timestamp, user_id, final_members
            FROM running_tasks
        ''')
        
        rows = cursor.fetchall()
        conn.close()
        
        return {row['task_id']: dict(row) for row in rows}
    
    def add_running_task(self, task_id: str, task_info: Dict[str, Any]):
        """添加运行中任务"""
        conn = self._get_connection()
        cursor = conn.cursor()
        
        cursor.execute('''
            INSERT OR REPLACE INTO running_tasks 
            (task_id, job_id, prefix, stage, status, timestamp, user_id, final_members)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?)
        ''', (
            task_id,
            task_info.get('job_id', ''),
            task_info.get('prefix', ''),
            task_info.get('stage', ''),
            task_info.get('status', ''),
            task_info.get('timestamp', datetime.now().strftime('%Y-%m-%d %H:%M:%S')),
            task_info.get('user_id', ''),
            task_info.get('final_members')
        ))
        
        conn.commit()
        conn.close()
    
    def update_running_task(self, task_id: str, status: str, 
                           final_members: Optional[int] = None):
        """更新运行中任务状态"""
        conn = self._get_connection()
        cursor = conn.cursor()
        
        if final_members is not None:
            cursor.execute('''
                UPDATE running_tasks 
                SET status = ?, final_members = ?, updated_at = CURRENT_TIMESTAMP
                WHERE task_id = ?
            ''', (status, final_members, task_id))
        else:
            cursor.execute('''
                UPDATE running_tasks 
                SET status = ?, updated_at = CURRENT_TIMESTAMP
                WHERE task_id = ?
            ''', (status, task_id))
        
        conn.commit()
        conn.close()
    
    def remove_running_task(self, task_id: str):
        """移除运行中任务"""
        conn = self._get_connection()
        cursor = conn.cursor()
        
        cursor.execute('''
            DELETE FROM running_tasks WHERE task_id = ?
        ''', (task_id,))
        
        conn.commit()
        conn.close()
    
    # ==================== 统计和清理 ====================
    
    def get_user_stats(self, user_id: str) -> Dict[str, int]:
        """获取用户任务统计"""
        conn = self._get_connection()
        cursor = conn.cursor()
        
        cursor.execute('''
            SELECT status, COUNT(*) as count
            FROM user_sessions
            WHERE user_id = ?
            GROUP BY status
        ''', (user_id,))
        
        rows = cursor.fetchall()
        conn.close()
        
        stats = {'running': 0, 'completed': 0, 'failed': 0}
        for row in rows:
            stats[row['status']] = row['count']
        return stats
    
    def cleanup_old_sessions(self, days: int = 30):
        """清理旧会话数据"""
        conn = self._get_connection()
        cursor = conn.cursor()
        
        cursor.execute('''
            DELETE FROM user_sessions 
            WHERE created_at < datetime('now', '-{} days')
        '''.format(days))
        
        deleted = cursor.rowcount
        conn.commit()
        conn.close()
        
        return deleted
    
    def close(self):
        """关闭存储（SQLite 不需要显式关闭，但保留接口一致性）"""
        pass


# 全局存储实例
_storage_instance: Optional[SessionStorage] = None


def get_storage() -> SessionStorage:
    """获取全局存储实例"""
    global _storage_instance
    if _storage_instance is None:
        _storage_instance = SessionStorage()
    return _storage_instance


def init_storage(db_path: str = DEFAULT_DB_PATH) -> SessionStorage:
    """初始化存储并返回实例"""
    global _storage_instance
    _storage_instance = SessionStorage(db_path)
    return _storage_instance