import aiosqlite
import json
from typing import Any, Dict, Optional
from aiogram.fsm.storage.base import BaseStorage, StateType, StorageKey

class SQLiteStorage(BaseStorage):
    """
    Simple SQLite-based FSM storage for aiogram 3.
    """
    def __init__(self, db_path: str = "fsm_storage.sqlite"):
        self.db_path = db_path
        self._conn = None

    async def _get_conn(self):
        if self._conn is None:
            self._conn = await aiosqlite.connect(self.db_path)
            await self._conn.execute('''
                CREATE TABLE IF NOT EXISTS fsm_data (
                    bot_id INTEGER,
                    chat_id INTEGER,
                    user_id INTEGER,
                    destiny TEXT,
                    state TEXT,
                    data TEXT,
                    PRIMARY KEY (bot_id, chat_id, user_id, destiny)
                )
            ''')
            await self._conn.commit()
        return self._conn

    async def close(self) -> None:
        if self._conn:
            await self._conn.close()
            self._conn = None

    async def set_state(self, key: StorageKey, state: StateType = None) -> None:
        conn = await self._get_conn()
        state_str = state.state if hasattr(state, 'state') else state
        
        # We must insert or update the state
        await conn.execute('''
            INSERT INTO fsm_data (bot_id, chat_id, user_id, destiny, state, data)
            VALUES (?, ?, ?, ?, ?, '{}')
            ON CONFLICT(bot_id, chat_id, user_id, destiny) DO UPDATE SET state = excluded.state
        ''', (key.bot_id, key.chat_id, key.user_id, key.destiny, state_str))
        await conn.commit()

    async def get_state(self, key: StorageKey) -> Optional[str]:
        conn = await self._get_conn()
        async with conn.execute('''
            SELECT state FROM fsm_data 
            WHERE bot_id=? AND chat_id=? AND user_id=? AND destiny=?
        ''', (key.bot_id, key.chat_id, key.user_id, key.destiny)) as cursor:
            row = await cursor.fetchone()
            return row[0] if row else None

    async def set_data(self, key: StorageKey, data: Dict[str, Any]) -> None:
        conn = await self._get_conn()
        data_str = json.dumps(data)
        
        await conn.execute('''
            INSERT INTO fsm_data (bot_id, chat_id, user_id, destiny, state, data)
            VALUES (?, ?, ?, ?, NULL, ?)
            ON CONFLICT(bot_id, chat_id, user_id, destiny) DO UPDATE SET data = excluded.data
        ''', (key.bot_id, key.chat_id, key.user_id, key.destiny, data_str))
        await conn.commit()

    async def get_data(self, key: StorageKey) -> Dict[str, Any]:
        conn = await self._get_conn()
        async with conn.execute('''
            SELECT data FROM fsm_data 
            WHERE bot_id=? AND chat_id=? AND user_id=? AND destiny=?
        ''', (key.bot_id, key.chat_id, key.user_id, key.destiny)) as cursor:
            row = await cursor.fetchone()
            if row and row[0]:
                try:
                    return json.loads(row[0])
                except:
                    return {}
            return {}