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.execute("""
            CREATE TABLE IF NOT EXISTS memory_entries (
                id TEXT PRIMARY KEY,
                scope TEXT NOT NULL,
                project TEXT,
                type TEXT NOT NULL,
                title TEXT NOT NULL,
                body TEXT NOT NULL,
                tags TEXT,
                status TEXT NOT NULL DEFAULT 'active',
                superseded_by TEXT,
                created_by TEXT NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """)
        await db.execute("""
            CREATE VIRTUAL TABLE IF NOT EXISTS memory_entries_fts USING fts5(
                title, body, tags, content='memory_entries', content_rowid='rowid'
            )
        """)
        await db.execute("""
            CREATE TRIGGER IF NOT EXISTS memory_entries_ai AFTER INSERT ON memory_entries BEGIN
                INSERT INTO memory_entries_fts(rowid, title, body, tags) VALUES (new.rowid, new.title, new.body, new.tags);
            END
        """)
        await db.execute("""
            CREATE TRIGGER IF NOT EXISTS memory_entries_ad AFTER DELETE ON memory_entries BEGIN
                INSERT INTO memory_entries_fts(memory_entries_fts, rowid, title, body, tags) VALUES('delete', old.rowid, old.title, old.body, old.tags);
            END
        """)
        await db.execute("""
            CREATE TRIGGER IF NOT EXISTS memory_entries_au AFTER UPDATE ON memory_entries BEGIN
                INSERT INTO memory_entries_fts(memory_entries_fts, rowid, title, body, tags) VALUES('delete', old.rowid, old.title, old.body, old.tags);
                INSERT INTO memory_entries_fts(rowid, title, body, tags) VALUES (new.rowid, new.title, new.body, new.tags);
            END
        """)
        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]


_TYPE_ABBR = {
    "decision": "dec",
    "constraint": "cst",
    "best-practice": "bp",
    "common-error": "err",
    "do-not-use": "dnu",
}

async def _next_entry_id(db, scope: str, type_: str) -> str:
    abbr = _TYPE_ABBR.get(type_, type_)
    prefix = f"{scope}-{abbr}-"
    async with db.execute(
        "SELECT id FROM memory_entries WHERE scope=? AND type=? ORDER BY created_at DESC, id DESC LIMIT 1",
        (scope, type_)
    ) as cur:
        row = await cur.fetchone()
    n = 1
    if row:
        try:
            n = int(row[0].rsplit("-", 1)[-1]) + 1
        except ValueError:
            n = 1
    return f"{prefix}{n:04d}"

async def add_memory_entry(scope: str, type_: str, title: str, body: str, created_by: str,
                            project: Optional[str] = None, tags: Optional[List[str]] = None) -> Dict[str, Any]:
    tags_json = json.dumps(tags or [])
    async with aiosqlite.connect(DB_PATH) as db:
        entry_id = await _next_entry_id(db, scope, type_)
        await db.execute(
            """INSERT INTO memory_entries (id, scope, project, type, title, body, tags, created_by)
               VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
            (entry_id, scope, project, type_, title, body, tags_json, created_by)
        )
        await db.commit()
    return {"id": entry_id, "scope": scope, "project": project, "type": type_,
            "title": title, "body": body, "tags": tags or [], "status": "active", "created_by": created_by}

async def list_memory_entries(scope: str, project: Optional[str] = None,
                               type_: Optional[str] = None, status: str = "active") -> List[Dict[str, Any]]:
    query = ("SELECT id, scope, project, type, title, body, tags, status, superseded_by, "
             "created_by, created_at, updated_at FROM memory_entries WHERE scope=?")
    params: List[Any] = [scope]
    if project is not None:
        query += " AND project=?"
        params.append(project)
    if type_ is not None:
        query += " AND type=?"
        params.append(type_)
    if status:
        query += " AND status=?"
        params.append(status)
    query += " ORDER BY created_at DESC"
    async with aiosqlite.connect(DB_PATH) as db:
        db.row_factory = aiosqlite.Row
        async with db.execute(query, params) as cursor:
            rows = await cursor.fetchall()
            result = []
            for row in rows:
                d = dict(row)
                d["tags"] = json.loads(d["tags"]) if d["tags"] else []
                result.append(d)
            return result

async def search_memory_entries(query_text: str, scope: Optional[str] = None) -> List[Dict[str, Any]]:
    sql = ("SELECT me.id, me.scope, me.project, me.type, me.title, me.body, me.tags, "
           "me.status, me.created_by, me.created_at "
           "FROM memory_entries_fts JOIN memory_entries me ON me.rowid = memory_entries_fts.rowid "
           "WHERE memory_entries_fts MATCH ? AND me.status = 'active'")
    params: List[Any] = [query_text]
    if scope:
        sql += " AND me.scope = ?"
        params.append(scope)
    sql += " ORDER BY rank"
    async with aiosqlite.connect(DB_PATH) as db:
        db.row_factory = aiosqlite.Row
        async with db.execute(sql, params) as cursor:
            rows = await cursor.fetchall()
            result = []
            for row in rows:
                d = dict(row)
                d["tags"] = json.loads(d["tags"]) if d["tags"] else []
                result.append(d)
            return result

async def supersede_memory_entry(entry_id: str, superseded_by: str):
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute(
            "UPDATE memory_entries SET status='superseded', superseded_by=?, updated_at=CURRENT_TIMESTAMP WHERE id=?",
            (superseded_by, entry_id)
        )
        await db.commit()

async def archive_memory_entry(entry_id: str):
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute(
            "UPDATE memory_entries SET status='archived', updated_at=CURRENT_TIMESTAMP WHERE id=?",
            (entry_id,)
        )
        await db.commit()