import json
import re
from typing import Optional

from sqlalchemy import Engine, create_engine, inspect, text
from sqlalchemy.engine import Result, make_url
from sqlalchemy.exc import ArgumentError


class SqlClient:
    def __init__(self, db_url: str):
        self._engine = self.get_engine(db_url)

    def run(self, sql: str) -> tuple[Optional[list[dict]], Optional[str]]:
        ok, err = self.validate_sql(sql)
        if not ok:
            return None, err

        with self._engine.connect() as conn:
            result = conn.execute(text(sql))
            return [dict(row._mapping) for row in result], None

    def get_engine(self, url: str) -> Engine:
        try:
            engine = create_engine(url)

            if not self.test_connection(engine):
                raise ValueError("Unable to connect to the database")

            return engine
        except Exception:
            raise ValueError("Invalid database URL or unable to connect")

    def test_connection(self, engine: Engine) -> bool:
        try:
            with engine.connect() as conn:
                conn.execute(text("SELECT 1"))
            return True
        except Exception:
            return False

    def get_dialect(self) -> str:
        dialect = self._engine.dialect.name
        return dialect

    def get_schema_database(self) -> str:
        inspector = inspect(self._engine)
        schema = {}

        for table in inspector.get_table_names():
            schema[table] = {
                "columns": self._get_columns(inspector, table),
                "primary_key": self._get_primary_key(inspector, table),
                "foreign_keys": self._get_foreign_keys(inspector, table),
                "sample_data": self._get_sample_data(table),
            }

        return json.dumps(schema, ensure_ascii=False)

    def validate_sql(self, sql: str) -> tuple[bool, Optional[str]]:
        if not self.is_readonly_sql(sql):
            return False, "Only read-only SELECT queries are allowed."

        try:
            with self._engine.connect() as conn:
                conn.execute(text(f"EXPLAIN {sql}"))
            return True, None
        except Exception as e:
            return False, str(e)

    @staticmethod
    def validate_connection_string(
        connection_string: str,
    ) -> tuple[bool, Optional[str]]:
        """Check if a connection string is valid.

        Step 1: Validate URL syntax via make_url() (no network call).
        Step 2: Attempt a real connection to verify credentials/reachability.

        Returns:
            (True, None) if the connection is successful.
            (False, error_message) if syntax or connection fails.
        """
        # Step 1: syntax check
        try:
            make_url(connection_string)
        except ArgumentError as e:
            return False, f"Invalid URL format: {e}"

        # Step 2: real connection check
        try:
            engine = create_engine(connection_string)
            with engine.connect() as conn:
                conn.execute(text("SELECT 1"))
            engine.dispose()
            return True, None
        except Exception as e:
            return False, f"Connection failed: {e}"

    @staticmethod
    def extract_sql(text_response: str) -> str:
        patterns = [
            r"```sql\s*(.*?)```",
            r"```(.*?)```",
            r"\bWITH\b.*?;",
            r"\bSELECT\b.*?;",
        ]

        for p in patterns:
            matches = re.findall(p, text_response, re.DOTALL | re.IGNORECASE)
            if matches:
                return matches[-1].strip()

        return text_response.strip()

    @staticmethod
    def is_readonly_sql(sql: str) -> bool:
        sql_lower = sql.strip().lower()
        forbidden = (
            "insert ",
            "update ",
            "delete ",
            "drop ",
            "alter ",
            "truncate ",
            "create ",
        )
        return not any(word in sql_lower for word in forbidden)

    def _get_columns(self, inspector, table: str) -> list[dict]:
        return [
            {
                "name": col["name"],
                "type": str(col["type"]),
                "nullable": col["nullable"],
                "default": str(col["default"]) if col.get("default") else None,
            }
            for col in inspector.get_columns(table)
        ]

    def _get_primary_key(self, inspector, table: str) -> list[str]:
        pk = inspector.get_pk_constraint(table)
        return pk.get("constrained_columns") or []

    def _get_foreign_keys(self, inspector, table: str) -> list[dict]:
        return [
            {
                "columns": fk["constrained_columns"],
                "referred_table": fk["referred_table"],
                "referred_columns": fk["referred_columns"],
            }
            for fk in inspector.get_foreign_keys(table)
        ]

    def _get_sample_data(self, table: str) -> list[dict]:
        try:
            with self._engine.connect() as conn:
                result: Result = conn.execute(text(f'SELECT * FROM "{table}" LIMIT 5'))
                return [dict(row._mapping) for row in result]
        except Exception:
            return []
