from typing import Any

from cryptography.hazmat.primitives.serialization import (  # noqa: I001
    Encoding,
    load_pem_private_key,
    NoEncryption,
    PrivateFormat,
)
import snowflake.connector

from .models import ExecuteResult
from .secrets import get_secret


class SnowflakeConnection:
    def __init__(self, key_secret: str, **connect_kwargs: Any) -> None:
        """Connect to Snowflake using a private key from Secrets Manager.

        key_secret is the name of a Secrets Manager secret containing the private key
        as a PEM string. Bare base64 keys without -----BEGIN/END----- headers are
        wrapped automatically. All other kwargs are passed directly to
        snowflake.connector.connect().
        """
        raw_key = get_secret(key_secret)
        if not isinstance(raw_key, str):
            raise ValueError(
                f"Secret '{key_secret}' must be a plain string, got {type(raw_key).__name__}"
            )
        trimmed = raw_key.strip()
        if not trimmed.startswith('-----'):
            trimmed = (
                f'-----BEGIN PRIVATE KEY-----\n{trimmed}\n-----END PRIVATE KEY-----'
            )
        private_key = load_pem_private_key(trimmed.encode(), password=None)
        private_key_der = private_key.private_bytes(
            encoding=Encoding.DER,
            format=PrivateFormat.PKCS8,
            encryption_algorithm=NoEncryption(),
        )
        self._conn = snowflake.connector.connect(
            **connect_kwargs, private_key=private_key_der
        )

    def __enter__(self) -> 'SnowflakeConnection':
        return self

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

    def fetchall(
        self, sql: str, params: tuple[Any, ...] | None = None
    ) -> list[dict[str, Any]]:
        with self._conn.cursor(snowflake.connector.DictCursor) as cursor:
            cursor.execute(sql, params)
            return cursor.fetchall()

    def fetchone(
        self, sql: str, params: tuple[Any, ...] | None = None
    ) -> dict[str, Any] | None:
        with self._conn.cursor(snowflake.connector.DictCursor) as cursor:
            cursor.execute(sql, params)
            return cursor.fetchone()

    def execute(self, sql: str, params: tuple[Any, ...] | None = None) -> ExecuteResult:
        with self._conn.cursor(snowflake.connector.DictCursor) as cursor:
            cursor.execute(sql, params)
            rows = cursor.fetchall() if cursor.description is not None else []
            return ExecuteResult(rowcount=cursor.rowcount or 0, rows=rows)

    def close(self) -> None:
        self._conn.close()