import base64
import os
from typing import List, Dict
from contextlib import contextmanager

import snowflake.connector


class SnowflakeClient:
    def __init__(self):
        self.account = os.getenv("SNOWFLAKE_ACCOUNT")
        self.user = os.getenv("SNOWFLAKE_USER")
        self.password = os.getenv("SNOWFLAKE_PASSWORD")
        self.private_key_raw = os.getenv("SNOWFLAKE_KEY")
        self.role = os.getenv("SNOWFLAKE_ROLE")
        self.database = os.getenv("SNOWFLAKE_DATABASE")
        self.warehouse = os.getenv("SNOWFLAKE_WAREHOUSE")
        self.schema = os.getenv("SNOWFLAKE_SCHEMA", "public")
        self.connection = None

    def connect(self):
        kwargs = {
            "account": self.account,
            "user": self.user,
            "database": self.database,
            "warehouse": self.warehouse,
            "schema": self.schema,
        }
        if self.role:
            kwargs["role"] = self.role

        if self.private_key_raw:
            kwargs["private_key"] = base64.b64decode(
                bytes(self.private_key_raw, encoding="utf-8")
            )
        elif self.password:
            kwargs["password"] = self.password
        else:
            raise ValueError("No Snowflake credentials: set SNOWFLAKE_KEY or SNOWFLAKE_PASSWORD")

        try:
            self.connection = snowflake.connector.connect(**kwargs)
        except Exception as e:
            raise ConnectionError(f"Failed to connect to Snowflake: {e}")

    def disconnect(self):
        if self.connection:
            self.connection.close()

    @contextmanager
    def cursor(self):
        if not self.connection:
            self.connect()
        cursor = self.connection.cursor()
        try:
            yield cursor
        finally:
            cursor.close()

    def query(self, sql: str, parameters: tuple = None) -> List[Dict]:
        with self.cursor() as cursor:
            cursor.execute(sql, parameters or ())
            columns = [desc[0] for desc in cursor.description] if cursor.description else []
            return [dict(zip(columns, row)) for row in cursor.fetchall()]

    def get_table_info(self, table: str) -> Dict:
        sql = f"DESCRIBE TABLE {table}"
        with self.cursor() as cursor:
            cursor.execute(sql)
            columns = [desc[0] for desc in cursor.description]
            rows = cursor.fetchall()
            return {
                "columns": columns,
                "schema": [dict(zip(columns, row)) for row in rows]
            }