import aiosqlite, random, string

DB_PATH = "maxservice.db"

def gen_key():
    return ''.join(random.choices(string.digits, k=5))

async def init_db():
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute("""CREATE TABLE IF NOT EXISTS users (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            access_key TEXT UNIQUE NOT NULL,
            telegram_chat_id TEXT,
            key_sent INTEGER DEFAULT 0,
            created_at DATETIME DEFAULT CURRENT_TIMESTAMP)""")
        await db.execute("""CREATE TABLE IF NOT EXISTS sessions (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            user_id INTEGER REFERENCES users(id),
            session_json TEXT,
            status TEXT DEFAULT 'pending',
            poll_interval INTEGER DEFAULT 300,
            mute_filter INTEGER DEFAULT 1,
            notify_telegram INTEGER DEFAULT 0,
            notify_webpush INTEGER DEFAULT 0,
            push_subscription TEXT,
            created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
            updated_at DATETIME DEFAULT CURRENT_TIMESTAMP)""")
        for col, defn in [
            ("access_key", "TEXT UNIQUE"), ("telegram_chat_id","TEXT"), ("key_sent","INTEGER DEFAULT 0"),
            ("notify_telegram","INTEGER DEFAULT 0"), ("notify_webpush","INTEGER DEFAULT 0"),
            ("push_subscription","TEXT"),
        ]:
            if col == "access_key":
                try: await db.execute(f"ALTER TABLE users ADD COLUMN {col} {defn}")
                except: pass
            else:
                try: await db.execute(f"ALTER TABLE users ADD COLUMN {col} {defn}")
                except: pass
                try: await db.execute(f"ALTER TABLE sessions ADD COLUMN {col} {defn}")
                except: pass
        await db.commit()

async def create_user() -> dict:
    # Генерируем уникальный 5-значный ключ
    async with aiosqlite.connect(DB_PATH) as db:
        for _ in range(100):
            key = gen_key()
            async with db.execute("SELECT id FROM users WHERE access_key=?", (key,)) as cur:
                if not await cur.fetchone():
                    break
        await db.execute("INSERT INTO users (access_key) VALUES (?)", (key,))
        await db.commit()
        db.row_factory = aiosqlite.Row
        async with db.execute("SELECT * FROM users WHERE access_key=?", (key,)) as cur:
            return dict(await cur.fetchone())

async def get_user_by_key(key: str) -> dict | None:
    async with aiosqlite.connect(DB_PATH) as db:
        db.row_factory = aiosqlite.Row
        async with db.execute("SELECT * FROM users WHERE access_key=?", (key,)) as cur:
            row = await cur.fetchone()
            return dict(row) if row else None

async def set_telegram(user_id: int, tg_id: str):
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute("UPDATE users SET telegram_chat_id=? WHERE id=?", (tg_id, user_id))
        await db.commit()

async def mark_key_sent(user_id: int):
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute("UPDATE users SET key_sent=1 WHERE id=?", (user_id,))
        await db.commit()

async def get_session(user_id: int) -> dict | None:
    async with aiosqlite.connect(DB_PATH) as db:
        db.row_factory = aiosqlite.Row
        async with db.execute("SELECT * FROM sessions WHERE user_id=? ORDER BY id DESC LIMIT 1", (user_id,)) as cur:
            row = await cur.fetchone()
            return dict(row) if row else None

async def save_session(user_id: int, session_json: str):
    # Добавляем таймаут, чтобы SQLite ждал освобождения базы, а не падал сразу
    async with aiosqlite.connect(DB_PATH, timeout=30.0) as db:
        # Принудительно проверяем существование сессии
        async with db.execute("SELECT id FROM sessions WHERE user_id=?", (user_id,)) as cur:
            row = await cur.fetchone()
        
        if row:
            await db.execute(
                "UPDATE sessions SET session_json=?, status='active', updated_at=CURRENT_TIMESTAMP WHERE user_id=?", 
                (session_json, user_id)
            )
        else:
            await db.execute(
                "INSERT INTO sessions (user_id, session_json, status) VALUES (?, ?, 'active')", 
                (user_id, session_json)
            )
        await db.commit()

async def set_session_status(user_id: int, status: str):
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute("UPDATE sessions SET status=?, updated_at=CURRENT_TIMESTAMP WHERE user_id=?", (status, user_id))
        await db.commit()

async def update_settings(user_id: int, poll_interval: int, mute_filter: int, notify_telegram: int, notify_webpush: int):
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute("UPDATE sessions SET poll_interval=?, mute_filter=?, notify_telegram=?, notify_webpush=? WHERE user_id=?",
            (poll_interval, mute_filter, notify_telegram, notify_webpush, user_id))
        await db.commit()

async def save_push_subscription(user_id: int, sub: str):
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute("UPDATE sessions SET push_subscription=? WHERE user_id=?", (sub, user_id))
        await db.commit()

async def get_all_active_sessions() -> list:
    async with aiosqlite.connect(DB_PATH) as db:
        db.row_factory = aiosqlite.Row
        async with db.execute("""SELECT s.*, u.access_key, u.telegram_chat_id, u.key_sent
            FROM sessions s JOIN users u ON u.id=s.user_id
            WHERE s.status='active' AND s.session_json IS NOT NULL""") as cur:
            return [dict(r) for r in await cur.fetchall()]