import aiosqlite
import json
import logging
from datetime import datetime
from typing import List, Dict, Any, Optional

DB_PATH = "data/context_hub.db"

logger = logging.getLogger(__name__)

async def init_db():
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute("""
            CREATE TABLE IF NOT EXISTS rules (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                scope TEXT UNIQUE NOT NULL,
                content TEXT NOT NULL
            )
        """)
        await db.execute("""
            CREATE TABLE IF NOT EXISTS api_keys (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                agent_name TEXT UNIQUE NOT NULL,
                api_key TEXT UNIQUE NOT NULL,
                last_used TIMESTAMP
            )
        """)
        await db.execute("""
            CREATE TABLE IF NOT EXISTS audit_log (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                agent_name TEXT,
                scope TEXT,
                timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                ip_address TEXT
            )
        """)
        await db.execute("""
            CREATE TABLE IF NOT EXISTS ports (
                port INTEGER PRIMARY KEY,
                service_name TEXT NOT NULL,
                description TEXT
            )
        """)
        await db.commit()

async def get_rule(scope: str) -> Optional[Dict[str, Any]]:
    async with aiosqlite.connect(DB_PATH) as db:
        async with db.execute("SELECT content FROM rules WHERE scope = ?", (scope,)) as cursor:
            row = await cursor.fetchone()
            if row:
                return json.loads(row[0])
            return None

async def set_rule(scope: str, content: Dict[str, Any]):
    content_str = json.dumps(content)
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute(
            "INSERT INTO rules (scope, content) VALUES (?, ?) ON CONFLICT(scope) DO UPDATE SET content=?",
            (scope, content_str, content_str)
        )
        await db.commit()

async def get_all_rules() -> Dict[str, Any]:
    async with aiosqlite.connect(DB_PATH) as db:
        async with db.execute("SELECT scope, content FROM rules") as cursor:
            rows = await cursor.fetchall()
            return {row[0]: json.loads(row[1]) for row in rows}

async def get_all_ports() -> List[Dict[str, Any]]:
    async with aiosqlite.connect(DB_PATH) as db:
        db.row_factory = aiosqlite.Row
        async with db.execute("SELECT port, service_name, description FROM ports") as cursor:
            rows = await cursor.fetchall()
            return [dict(row) for row in rows]

async def add_port(port: int, service_name: str, description: str = ""):
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute(
            "INSERT INTO ports (port, service_name, description) VALUES (?, ?, ?) ON CONFLICT(port) DO UPDATE SET service_name=?, description=?",
            (port, service_name, description, service_name, description)
        )
        await db.commit()

async def get_port(port: int) -> Optional[Dict[str, Any]]:
    async with aiosqlite.connect(DB_PATH) as db:
        db.row_factory = aiosqlite.Row
        async with db.execute("SELECT port, service_name, description FROM ports WHERE port = ?", (port,)) as cursor:
            row = await cursor.fetchone()
            if row:
                return dict(row)
            return None

async def add_api_key(agent_name: str, api_key: str):
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute(
            "INSERT INTO api_keys (agent_name, api_key) VALUES (?, ?) ON CONFLICT(agent_name) DO UPDATE SET api_key=?",
            (agent_name, api_key, api_key)
        )
        await db.commit()

async def get_agent_by_key(api_key: str) -> Optional[str]:
    async with aiosqlite.connect(DB_PATH) as db:
        async with db.execute("SELECT agent_name FROM api_keys WHERE api_key = ?", (api_key,)) as cursor:
            row = await cursor.fetchone()
            if row:
                return row[0]
            return None

async def update_key_last_used(api_key: str):
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute("UPDATE api_keys SET last_used = CURRENT_TIMESTAMP WHERE api_key = ?", (api_key,))
        await db.commit()

async def get_all_keys() -> List[Dict[str, Any]]:
    async with aiosqlite.connect(DB_PATH) as db:
        db.row_factory = aiosqlite.Row
        async with db.execute("SELECT id, agent_name, api_key, last_used FROM api_keys") as cursor:
            rows = await cursor.fetchall()
            return [dict(row) for row in rows]

async def add_audit_log(agent_name: str, scope: str, ip_address: str):
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute(
            "INSERT INTO audit_log (agent_name, scope, ip_address) VALUES (?, ?, ?)",
            (agent_name, scope, ip_address)
        )
        await db.commit()

async def get_audit_logs(limit: int = 50) -> List[Dict[str, Any]]:
    async with aiosqlite.connect(DB_PATH) as db:
        db.row_factory = aiosqlite.Row
        async with db.execute("SELECT id, agent_name, scope, timestamp, ip_address FROM audit_log ORDER BY timestamp DESC LIMIT ?", (limit,)) as cursor:
            rows = await cursor.fetchall()
            return [dict(row) for row in rows]