import json
from typing import Any, Dict, List, Optional

from agno.tools import Toolkit
from agno.utils.log import log_debug, logger

try:
    from sqlalchemy import Engine, create_engine
    from sqlalchemy.inspection import inspect
    from sqlalchemy.orm import Session, sessionmaker
    from sqlalchemy.sql.expression import text
except ImportError:
    raise ImportError("`sqlalchemy` not installed")


class SQLTools(Toolkit):
    def __init__(
        self,
        db_url: Optional[str] = None,
        db_engine: Optional[Engine] = None,
        user: Optional[str] = None,
        password: Optional[str] = None,
        host: Optional[str] = None,
        port: Optional[int] = None,
        schema: Optional[str] = None,
        dialect: Optional[str] = None,
        tables: Optional[Dict[str, Any]] = None,
        enable_list_tables: bool = True,
        enable_describe_table: bool = True,
        enable_run_sql_query: bool = True,
        all: bool = False,
        **kwargs,
    ):
        # Get the database engine
        _engine: Optional[Engine] = db_engine
        if _engine is None and db_url is not None:
            _engine = create_engine(db_url)
        elif user and password and host and port and dialect:
            if schema is not None:
                _engine = create_engine(f"{dialect}://{user}:{password}@{host}:{port}/{schema}")
            else:
                _engine = create_engine(f"{dialect}://{user}:{password}@{host}:{port}")

        if _engine is None:
            raise ValueError("Could not build the database connection")

        # Database connection
        self.db_engine: Engine = _engine
        self.Session: sessionmaker[Session] = sessionmaker(bind=self.db_engine)

        self.schema = schema

        # Tables this toolkit can access
        self.tables: Optional[Dict[str, Any]] = tables

        tools: List[Any] = []
        if enable_list_tables or all:
            tools.append(self.list_tables)
        if enable_describe_table or all:
            tools.append(self.describe_table)
        if enable_run_sql_query or all:
            tools.append(self.run_sql_query)

        super().__init__(name="sql_tools", tools=tools, **kwargs)

    def list_tables(self) -> str:
        """Use this function to get a list of table names in the database with row counts.

        Returns:
            str: list of tables and row counts in the database.
        """
        if self.tables is not None:
            return json.dumps(self.tables)

        try:
            log_debug("Listing tables in the database")
            inspector = inspect(self.db_engine)

            # → Get all table names
            if self.schema:
                table_names = inspector.get_table_names(schema=self.schema)
            else:
                table_names = inspector.get_table_names()

            log_debug(f"table_names: {table_names}")

            # → Count rows for each table
            table_info = {}
            with self.Session() as session:
                for table in table_names:
                    try:
                        # Safe generic SQL for all dialects
                        count_query = f"SELECT COUNT(*) AS count FROM {table}"
                        result = session.execute(text(count_query)).scalar()
                    except Exception as e:
                        result = f"Error counting rows: {e}"

                    table_info[table] = {"row_count": result}

            return json.dumps(table_info, indent=2)

        except Exception as e:
            logger.error(f"Error getting tables: {e}")
            return f"Error getting tables: {e}"


    def describe_table(self, table_name: str) -> str:
        """Use this function to describe a table in detail.

        Args:
            table_name (str): The name of the table to get schema for.

        Returns:
            str: JSON string with table schema + keys + constraints.
        """
        try:
            log_debug(f"Describing table: {table_name}")
            inspector = inspect(self.db_engine)

            # → Columns
            columns = inspector.get_columns(table_name, schema=self.schema)
            column_details = []
            for col in columns:
                column_details.append({
                    "name": col.get("name"),
                    "type": str(col.get("type")),
                    "nullable": col.get("nullable"),
                    "default": str(col.get("default")),
                    "primary_key": col.get("primary_key", False)
                })

            # → Primary Keys
            primary_keys = inspector.get_primary_keys(table_name, schema=self.schema)

            # → Foreign Keys
            fk_raw = inspector.get_foreign_keys(table_name, schema=self.schema)
            foreign_keys = []
            for fk in fk_raw:
                foreign_keys.append({
                    "constrained_columns": fk.get("constrained_columns"),
                    "referred_table": fk.get("referred_table"),
                    "referred_columns": fk.get("referred_columns"),
                    "name": fk.get("name"),
                })

            # → Unique Constraints
            uc_raw = inspector.get_unique_constraints(table_name, schema=self.schema)
            unique_constraints = [
                {
                    "name": uc.get("name"),
                    "column_names": uc.get("column_names")
                }
                for uc in uc_raw
            ]

            result = {
                "table": table_name,
                "columns": column_details,
                "primary_keys": primary_keys,
                "foreign_keys": foreign_keys,
                "unique_constraints": unique_constraints,
            }

            return json.dumps(result, indent=2)

        except Exception as e:
            logger.error(f"Error getting table schema: {e}")
            return f"Error getting table schema: {e}"

    def run_sql_query(self, query: str, limit: Optional[int] = 10) -> str:
        """Use this function to run a SQL query and return the result.

        Args:
            query (str): The query to run.
            limit (int, optional): The number of rows to return. Defaults to 10. Use `None` to show all results.
        Returns:
            str: Result of the SQL query.
        Notes:
            - The result may be empty if the query does not return any data.
        """

        try:
            return json.dumps(self.run_sql(sql=query, limit=limit), default=str)
        except Exception as e:
            logger.error(f"Error running query: {e}")
            return f"Error running query: {e}"

    def run_sql(self, sql: str, limit: Optional[int] = None) -> List[dict]:
        """Internal function to run a sql query.

        Args:
            sql (str): The sql query to run.
            limit (int, optional): The number of rows to return. Defaults to None.

        Returns:
            List[dict]: The result of the query.
        """
        log_debug(f"Running sql |\n{sql}")

        with self.Session() as sess, sess.begin():
            result = sess.execute(text(sql))

            # Check if the operation has returned rows.
            try:
                if limit:
                    rows = result.fetchmany(limit)
                else:
                    rows = result.fetchall()
                return [row._asdict() for row in rows]
            except Exception as e:
                logger.error(f"Error while executing SQL: {e}")
                return []
