import sqlite3
import json
from datetime import datetime
from pathlib import Path
import streamlit as st
from typing import Dict, Any, List, Optional, Tuple


class SessionDatabase:
    """会话数据库管理"""
    
    def __init__(self, db_path: str = None):
        if db_path is None:
            # 默认数据库路径在项目根目录下的data文件夹
            project_root = Path(__file__).parent.parent.parent
            data_dir = project_root / "data"
            data_dir.mkdir(exist_ok=True)
            self.db_path = data_dir / "sessions.db"
        else:
            self.db_path = Path(db_path)
        
        self.init_database()
    
    def init_database(self):
        """初始化数据库表"""
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS sessions (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    name TEXT NOT NULL UNIQUE,
                    created_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    session_data TEXT NOT NULL
                )
            ''')
            conn.commit()
    
    def save_session(self, session_name: str, session_data: Dict[str, Any], force_overwrite: bool = False) -> tuple[bool, bool]:
        """保存会话状态
        
        Returns:
            tuple[bool, bool]: (是否成功, 是否存在重名)
        """
        try:
            # 序列化session_data，处理不可序列化的对象
            serialized_data = self._serialize_session_data(session_data)
            
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                
                # 检查会话名是否已存在
                cursor.execute('SELECT id FROM sessions WHERE name = ?', (session_name,))
                existing = cursor.fetchone()
                
                if existing and not force_overwrite:
                    # 存在重名且不强制覆盖
                    return False, True
                elif existing:
                    # 更新现有会话
                    cursor.execute('''
                        UPDATE sessions 
                        SET session_data = ?, updated_time = CURRENT_TIMESTAMP 
                        WHERE name = ?
                    ''', (serialized_data, session_name))
                else:
                    # 创建新会话
                    cursor.execute('''
                        INSERT INTO sessions (name, session_data) 
                        VALUES (?, ?)
                    ''', (session_name, serialized_data))
                
                conn.commit()
                return True, bool(existing)
                
        except Exception as e:
            st.error(f"保存会话失败: {str(e)}")
            return False, False
    
    def load_session(self, session_name: str) -> Optional[Dict[str, Any]]:
        """加载会话状态"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute('''
                    SELECT session_data FROM sessions 
                    WHERE name = ?
                ''', (session_name,))
                
                result = cursor.fetchone()
                if result:
                    return self._deserialize_session_data(result[0])
                return None
                
        except Exception as e:
            st.error(f"加载会话失败: {str(e)}")
            return None
    
    def get_all_sessions(self) -> List[Tuple[str, datetime]]:
        """获取所有会话列表"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute('''
                    SELECT name, updated_time 
                    FROM sessions 
                    ORDER BY updated_time DESC
                ''')
                
                sessions = []
                for row in cursor.fetchall():
                    name, updated_time = row
                    # 解析时间戳
                    dt = datetime.fromisoformat(updated_time.replace('Z', '+00:00'))
                    sessions.append((name, dt))
                
                return sessions
                
        except Exception as e:
            st.error(f"获取会话列表失败: {str(e)}")
            return []
    
    def delete_session(self, session_name: str) -> bool:
        """删除会话"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute('DELETE FROM sessions WHERE name = ?', (session_name,))
                conn.commit()
                return cursor.rowcount > 0
                
        except Exception as e:
            st.error(f"删除会话失败: {str(e)}")
            return False
    
    def rename_session(self, old_name: str, new_name: str) -> bool:
        """重命名会话"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                
                # 检查新名称是否已存在
                cursor.execute('SELECT id FROM sessions WHERE name = ?', (new_name,))
                if cursor.fetchone():
                    st.error(f"会话名称 '{new_name}' 已存在")
                    return False
                
                # 重命名
                cursor.execute('''
                    UPDATE sessions 
                    SET name = ?, updated_time = CURRENT_TIMESTAMP 
                    WHERE name = ?
                ''', (new_name, old_name))
                
                conn.commit()
                return cursor.rowcount > 0
                
        except Exception as e:
            st.error(f"重命名会话失败: {str(e)}")
            return False
    
    def get_latest_session(self) -> Optional[str]:
        """获取最新的会话名称"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute('''
                    SELECT name FROM sessions 
                    ORDER BY updated_time DESC 
                    LIMIT 1
                ''')
                
                result = cursor.fetchone()
                return result[0] if result else None
                
        except Exception as e:
            st.error(f"获取最新会话失败: {str(e)}")
            return None
    
    def _serialize_session_data(self, session_data: Dict[str, Any]) -> str:
        """序列化会话数据，过滤掉不可序列化的对象"""
        serializable_data = {}
        
        for key, value in session_data.items():
            try:
                # 首先处理numpy数组
                if hasattr(value, 'dtype') and hasattr(value, 'tolist'):  # numpy array
                    serializable_data[key] = {
                        '_type': 'numpy_array',
                        '_data': value.tolist(),
                        '_dtype': str(value.dtype)
                    }
                    continue
                
                # 尝试序列化，如果失败则进入异常处理
                json.dumps(value)
                serializable_data[key] = value
                
            except (TypeError, ValueError):
                # 对于DataFrame等复杂对象，进行特殊处理
                if hasattr(value, 'to_dict'):
                    try:
                        # 处理pandas DataFrame
                        df_data = value.to_dict('records')
                        # 确保DataFrame数据中的numpy数组也被处理
                        processed_data = self._process_nested_data(df_data)
                        serializable_data[key] = {
                            '_type': 'dataframe',
                            '_data': processed_data
                        }
                    except Exception as e:
                        # 如果DataFrame处理失败，尝试保存基本信息
                        try:
                            serializable_data[key] = {
                                '_type': 'dataframe_info',
                                '_shape': getattr(value, 'shape', None),
                                '_columns': list(getattr(value, 'columns', [])),
                                '_error': str(e)
                            }
                        except:
                            continue
                elif hasattr(value, 'tolist'):  # numpy array fallback
                    try:
                        serializable_data[key] = {
                            '_type': 'numpy_array',
                            '_data': value.tolist()
                        }
                    except:
                        continue
                elif hasattr(value, '__dict__'):
                    try:
                        serializable_data[key] = {
                            '_type': 'object',
                            '_data': str(value)
                        }
                    except:
                        continue
                else:
                    # 跳过不可序列化的对象
                    continue
        
        return json.dumps(serializable_data, ensure_ascii=False, indent=2)
    
    def _process_nested_data(self, data):
        """递归处理嵌套数据中的numpy数组"""
        if isinstance(data, list):
            return [self._process_nested_data(item) for item in data]
        elif isinstance(data, dict):
            return {k: self._process_nested_data(v) for k, v in data.items()}
        elif hasattr(data, 'dtype') and hasattr(data, 'tolist'):  # numpy array
            return data.tolist()
        else:
            return data
    
    def _deserialize_session_data(self, serialized_data: str) -> Dict[str, Any]:
        """反序列化会话数据"""
        try:
            data = json.loads(serialized_data)
            
            # 处理特殊类型的反序列化
            for key, value in data.items():
                if isinstance(value, dict) and '_type' in value:
                    if value['_type'] == 'dataframe':
                        try:
                            import pandas as pd
                            data[key] = pd.DataFrame(value['_data'])
                        except ImportError:
                            data[key] = value['_data']
                    elif value['_type'] == 'numpy_array':
                        try:
                            import numpy as np
                            arr_data = value['_data']
                            if '_dtype' in value:
                                data[key] = np.array(arr_data, dtype=value['_dtype'])
                            else:
                                data[key] = np.array(arr_data)
                        except ImportError:
                            data[key] = value['_data']
                    elif value['_type'] == 'dataframe_info':
                        # 对于无法完全恢复的DataFrame，提供基本信息
                        data[key] = f"DataFrame(shape={value.get('_shape')}, columns={value.get('_columns')})"
                    elif value['_type'] == 'object':
                        data[key] = value['_data']
            
            return data
            
        except json.JSONDecodeError as e:
            st.error(f"反序列化会话数据失败: {str(e)}")
            return {}
    
    @staticmethod
    def generate_auto_session_name() -> str:
        """生成自动会话名称"""
        now = datetime.now()
        return f"auto_{now.strftime('%m%d%H%M')}"