import os
import mysql.connector
from typing import List, Dict
from contextlib import contextmanager


class MySQLClient:
    def __init__(self):
        self.host = os.getenv("MYSQL_HOST", "localhost")
        self.user = os.getenv("MYSQL_USER", "root")
        self.password = os.getenv("MYSQL_PASSWORD", "password")
        self.connection = None

    def connect(self):
        try:
            self.connection = mysql.connector.connect(
                host=self.host,
                user=self.user,
                password=self.password
            )
        except Exception as e:
            raise ConnectionError(f"Failed to connect to MySQL: {e}")

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

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

    def query(self, sql: str, parameters: tuple = None) -> List[Dict]:
        """Execute a SQL query and return results."""
        with self.cursor() as cursor:
            cursor.execute(sql, parameters or ())
            return cursor.fetchall()

    def get_table_schema(self, table: str) -> List[Dict]:
        """Get table schema information."""
        sql = f"DESCRIBE {table}"
        return self.query(sql)