from __future__ import annotations

import importlib
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple

from mockmesh.core.fallback import FallbackConfig
from mockmesh.core.logger import get_logger
from mockmesh.interceptors.sql.common import (
    MockSQLExecutor, MockQueryResult,
    load_sql_defaults, merge_overrides,
    get_connection_defaults, get_operation_defaults,
)
from mockmesh.plugins.base import BasePlugin

if TYPE_CHECKING:
    from mockmesh.response.engine import ResponseEngine
    from mockmesh.storage.manager import StorageManager

_log = get_logger("interceptors.sql")


class MockCursor:

    def __init__(self, executor: MockSQLExecutor,
                 defaults: Optional[dict] = None) -> None:
        self._executor   = executor
        self._defaults   = defaults
        self._result:    Optional[MockQueryResult] = None
        self.description: Optional[List[Tuple]] = None
        self.rowcount:   int = -1
        self.lastrowid:  Optional[int] = None
        self.arraysize:  int = 1

    def execute(self, sql: str, params: Any = None) -> None:
        self._result      = self._executor.execute(sql, params)
        self.description  = self._result.description
        self.rowcount     = self._result.rowcount
        self.lastrowid    = self._result.lastrowid

    def executemany(self, sql: str, seq_of_params: Any) -> None:
        count = 0
        for params in seq_of_params:
            result = self._executor.execute(sql, params)
            count += max(result.rowcount, 0)
        self.rowcount = count

    def callproc(self, procname: str, params: Any = None) -> None:
        op = get_operation_defaults("callproc", self._defaults)
        self._result = MockQueryResult(
            op.get("rows", []),
            rowcount=op.get("rowcount", 0),
        )
        self.rowcount = self._result.rowcount

    def fetchone(self) -> Optional[Tuple]:
        return self._result.fetchone() if self._result else None

    def fetchmany(self, size: Optional[int] = None) -> List[Tuple]:
        n = size if size is not None else self.arraysize
        return self._result.fetchmany(n) if self._result else []

    def fetchall(self) -> List[Tuple]:
        return self._result.fetchall() if self._result else []

    def __iter__(self):
        return iter(self._result.rows if self._result else [])

    def close(self) -> None:
        self._result = None

    def __enter__(self):
        return self

    def __exit__(self, *_: Any) -> None:
        self.close()

    def fetchone_dict(self) -> Optional[Dict[str, Any]]:
        row = self.fetchone()
        if row is None or self.description is None:
            return None
        return dict(zip([d[0] for d in self.description], row))

    def fetchall_dict(self) -> List[Dict[str, Any]]:
        if not self.description:
            return []
        cols = [d[0] for d in self.description]
        return [dict(zip(cols, row)) for row in self.fetchall()]


class MockConnection:

    def __init__(self, storage: "StorageManager",
                 database: Optional[str] = None,
                 driver: str = "unknown",
                 defaults: Optional[dict] = None) -> None:
        self._storage  = storage
        self._defaults = defaults
        conn = get_connection_defaults(defaults)
        self._database = database or conn.get("default_database", "mock")
        self._driver   = driver
        self._executor = MockSQLExecutor(storage, self._database, defaults)
        self._closed   = False
        self.autocommit = conn.get("autocommit", False)
        _log.debug("SQL mock connection opened",
                   extra={"mm_driver": driver, "mm_database": self._database})

    def cursor(self, *args: Any, **kwargs: Any) -> MockCursor:
        return MockCursor(self._executor, self._defaults)

    def commit(self) -> None:
        pass

    def rollback(self) -> None:
        pass

    def begin(self) -> None:
        pass

    def close(self) -> None:
        self._closed = True
        _log.debug("SQL mock connection closed",
                   extra={"mm_driver": self._driver, "mm_database": self._database})

    def __enter__(self) -> "MockConnection":
        return self

    def __exit__(self, *_: Any) -> None:
        self.close()

    def set_session(self, *args: Any, **kwargs: Any) -> None:
        pass

    def set_isolation_level(self, level: Any) -> None:
        pass

    def get_dsn_parameters(self) -> Dict[str, str]:
        conn = get_connection_defaults(self._defaults)
        return {
            "dbname": self._database,
            "host": conn.get("default_host", "localhost"),
        }

    async def __aenter__(self) -> "MockConnection":
        return self

    async def __aexit__(self, *_: Any) -> None:
        self.close()

    async def execute(self, sql: str, *args: Any) -> str:
        params = list(args) if args else None
        self._executor.execute(sql, params)
        return "OK"

    async def fetch(self, sql: str, *args: Any) -> List[Dict]:
        params = list(args) if args else None
        result = self._executor.execute(sql, params)
        if not result.description:
            return []
        cols = [d[0] for d in result.description]
        return [dict(zip(cols, row)) for row in result.fetchall()]

    async def fetchrow(self, sql: str, *args: Any) -> Optional[Dict]:
        rows = await self.fetch(sql, *args)
        return rows[0] if rows else None

    async def fetchval(self, sql: str, *args: Any) -> Any:
        row = await self.fetchrow(sql, *args)
        if row is None:
            return None
        return next(iter(row.values()))


_DRIVERS: List[Tuple[str, str, str]] = [
    ("psycopg2",        "connect",  "database"),
    ("psycopg",         "connect",  "dbname"),
    ("pymysql",         "connect",  "database"),
    ("mysql.connector", "connect",  "database"),
    ("aiomysql",        "connect",  "db"),
]


class SqlInterceptor(BasePlugin):

    name = "sql"

    def __init__(self, response_engine: "ResponseEngine") -> None:
        super().__init__(response_engine)
        self._originals: Dict[str, Any] = {}
        self._storage: Optional["StorageManager"] = None
        self._defaults: Optional[dict] = None

    def activate(
        self,
        config: dict,
        *,
        storage: Optional["StorageManager"] = None,
        provider_defaults: Optional[dict] = None,
        fallback_config: Optional[FallbackConfig] = None,
        custom_handlers: Optional[dict] = None,
        **_: Any,
    ) -> None:
        self._storage = storage
        base = load_sql_defaults()
        if provider_defaults and "sql" in provider_defaults:
            self._defaults = merge_overrides(base, provider_defaults["sql"])
        else:
            self._defaults = base
        self._patch_all()
        self._active = True
        _log.info("SQL interceptor activated",
                  extra={"mm_patched": [k for k in self._originals]})

    def deactivate(self) -> None:
        if not self._active:
            return
        self._unpatch_all()
        self._active = False
        self._defaults = None
        _log.info("SQL interceptor deactivated")

    def _patch_all(self) -> None:
        for module_name, attr, db_kwarg in _DRIVERS:
            try:
                mod = importlib.import_module(module_name)
            except ImportError:
                continue
            original = getattr(mod, attr, None)
            if original is None:
                continue
            self._originals[f"{module_name}.{attr}"] = original

            def _make_mock_connect(storage, database_kwarg, driver, defaults):
                def _mock_connect(*args, **kwargs):
                    conn = get_connection_defaults(defaults)
                    default_db = conn.get("default_database", "mock")
                    db = (kwargs.get(database_kwarg)
                          or kwargs.get("dbname")
                          or kwargs.get("db")
                          or (args[0] if args else default_db))
                    if isinstance(db, str) and "://" in db:
                        db = db.split("/")[-1].split("?")[0] or default_db
                    _log.info("SQL connect intercepted",
                              extra={"mm_driver": driver, "mm_db": db})
                    return MockConnection(storage, database=str(db),
                                          driver=driver, defaults=defaults)
                return _mock_connect

            setattr(mod, attr,
                    _make_mock_connect(self._storage, db_kwarg,
                                       module_name, self._defaults))
            _log.debug(f"Patched {module_name}.{attr}")

        try:
            import asyncpg  # type: ignore
            self._originals["asyncpg.connect"] = asyncpg.connect

            storage = self._storage
            defaults = self._defaults

            async def _mock_asyncpg_connect(*args, **kwargs):
                conn = get_connection_defaults(defaults)
                default_db = conn.get("default_database", "mock")
                db = kwargs.get("database") or kwargs.get("db") or default_db
                if args:
                    dsn = args[0]
                    if "://" in str(dsn):
                        db = str(dsn).split("/")[-1].split("?")[0] or default_db
                _log.info("asyncpg connect intercepted", extra={"mm_db": db})
                return MockConnection(storage, database=str(db),
                                      driver="asyncpg", defaults=defaults)

            asyncpg.connect = _mock_asyncpg_connect
            _log.debug("Patched asyncpg.connect")
        except ImportError:
            pass

    def _unpatch_all(self) -> None:
        for key, original in self._originals.items():
            module_name, attr = key.rsplit(".", 1)
            try:
                mod = importlib.import_module(module_name)
                setattr(mod, attr, original)
            except Exception:
                pass
        self._originals.clear()
