import json
import logging
from datetime import datetime
from typing import Any, Dict, List, Optional, Tuple
from urllib.parse import urlparse, urlunparse

from sqlalchemy import select, text, and_, or_, literal_column
from sqlalchemy.sql import Select
from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker

from .base import DBClient
from ..query import Query, WhereCondition, WhereGroup, QuerySet
from ..query.operators import Eq, Gt, Gte, Lt, Lte, In, Like, ILike, NotEq, IsNull, IsNotNull, RawExpr

logger = logging.getLogger(__name__)


def _serialize_pg_params(data: dict[str, Any]) -> dict[str, Any]:
    """Prepare values for asyncpg: JSON-encode dicts/lists, parse ISO timestamps.

    asyncpg requires ``datetime`` objects for TIMESTAMPTZ columns — plain
    ISO 8601 strings are rejected.
    """
    result: dict[str, Any] = {}
    for k, v in data.items():
        if isinstance(v, (dict, list)):
            result[k] = json.dumps(v)
        elif isinstance(v, str) and "T" in v:
            try:
                result[k] = datetime.fromisoformat(v)
            except (ValueError, TypeError):
                result[k] = v
        else:
            result[k] = v
    return result


class PostgresClient(DBClient):
    def __init__(self, url: str, echo: bool = False):
        self.engine = create_async_engine(
            url,
            pool_size=25,
            max_overflow=15,
            pool_pre_ping=True,
            echo=echo,
            connect_args={
                "server_settings": {"timezone": "UTC"},
            },
        )
        self.session_maker = async_sessionmaker(self.engine, expire_on_commit=False)

    async def connect(self):
        """Open a connection and verify the engine is ready."""
        async with self.engine.connect() as connection:
            await connection.scalar(text("SELECT 1"))

    async def disconnect(self):
        await self.engine.dispose()

    def table(self, name: str) -> "QueryBuilder":
        from ..query import QueryBuilder
        return QueryBuilder(self, name)

    def raw_query(self, query: Query) -> Tuple[str, Dict[str, Any]]:
        """Returns the final raw SQL query string and parameters"""
        stmt, params = self._build_select_statement(query, for_explain=True)
        compiled = stmt.compile(compile_kwargs={"literal_binds": True})
        sql_string = str(compiled)
        # SQLAlchemy emits LIMIT -1 for offset-only queries; PostgreSQL rejects it
        if " LIMIT -1" in sql_string:
            sql_string = sql_string.replace(" LIMIT -1", "")
        # Merge accumulated raw expression params
        if hasattr(query, '_params') and query._params:
            params = {**query._params, **params}
        return sql_string, params
    
    def debug_query(self, query: Query) -> str:
        """
        Returns a formatted debug string showing the SQL query.
        Useful for logging and debugging.
        
        Example:
            qb = db.table("users").where("age", ">", 18)
            print(db.debug_query(qb.query))
        """
        sql, params = self.raw_query(query)
        debug_str = f"SQL: {sql}"
        if params:
            debug_str += f"\nParams: {params}"
        return debug_str

    async def explain(self, query: Query, analyze: bool = False) -> str:
        """Return EXPLAIN output"""
        sql, _ = self.raw_query(query)
        explain_clause = "EXPLAIN ANALYZE" if analyze else "EXPLAIN"
        explain_sql = f"{explain_clause} {sql}"

        async with self.session_maker() as session:
            result = await session.execute(text(explain_sql))
            rows = result.fetchall()
            return "\n".join(str(row[0]) for row in rows)

    async def execute_raw(
        self,
        sql: str,
        params: dict[str, Any] | None = None,
        *,
        multi_statement: bool = False,
    ) -> Any:
        """Execute arbitrary raw SQL against PostgreSQL.

        Args:
            sql: Raw SQL string (may contain ``;``-delimited statements
                when ``multi_statement=True``).
            params: Optional parameter dict for bound parameters (``:name``).
                Parameters are only supported for single-statement mode.
            multi_statement: If True, the SQL may contain multiple statements
                separated by ``;``.  Uses a raw asyncpg connection to execute
                all statements in one round-trip.

        Returns:
            * ``list[dict[str, Any]]`` for SELECT-like statements.
            * ``int`` (rowcount) for INSERT/UPDATE/DELETE.
            * ``None`` for DDL like CREATE TABLE (multi-statement always
              returns ``None``).
        """
        if multi_statement:
            # Use a raw asyncpg connection — SQLAlchemy's text() does not
            # support multiple statements in a single execute() call.
            async with self.engine.connect() as conn:
                raw_conn = await conn.get_raw_connection()
                await raw_conn.driver_connection.execute(sql)
                await conn.commit()
            return None

        from sqlalchemy import text as sa_text

        async with self.session_maker() as session:
            result = await session.execute(sa_text(sql), params or {})
            if result.returns_rows:
                return [dict(row) for row in result.mappings()]
            await session.commit()
            return result.rowcount

    async def _execute(self, query: Query) -> QuerySet:
        sql, params = self.raw_query(query)
        async with self.session_maker() as session:
            result = await session.execute(text(sql), params)
            rows = [dict(row) for row in result.mappings()]
            return QuerySet(rows=rows, client=self, table=query.table, query=query)

    def _format_conditions(self, conditions: Dict[str, Any]) -> Tuple[str, Dict[str, Any]]:
        if not conditions:
            return "", {}

        clauses: List[str] = []
        params: Dict[str, Any] = {}
        for index, (field, value) in enumerate(conditions.items()):
            param_name = f"cond_{index}"
            if isinstance(value, list):
                in_params = [f":{param_name}_{i}" for i in range(len(value))]
                clauses.append(f"{field} IN ({', '.join(in_params)})")
                for i, v in enumerate(value):
                    params[f"{param_name}_{i}"] = v
            else:
                clauses.append(f"{field} = :{param_name}")
                params[param_name] = value
        return " AND ".join(clauses), params

    async def _insert_one(
        self,
        table: str,
        data: Dict[str, Any],
        returning: Optional[List[str]] = None
    ) -> Optional[Dict[str, Any]]:
        columns = ", ".join(data.keys())
        placeholders = ", ".join([f":{key}" for key in data.keys()])
        returning_clause = f" RETURNING {', '.join(returning)}" if returning else ""
        sql = f"INSERT INTO {table} ({columns}) VALUES ({placeholders}){returning_clause}"

        # Serialize dict/list → JSON and ISO strings → datetime for asyncpg
        safe_data = _serialize_pg_params(data)

        async with self.session_maker() as session:
            result = await session.execute(text(sql), safe_data)
            await session.commit()
            if returning:
                rows = result.mappings().all()
                return rows[0] if rows else None
            return None

    async def _insert_many(
        self,
        table: str,
        data: List[Dict[str, Any]],
        returning: Optional[List[str]] = None
    ) -> List[Dict[str, Any]]:
        if not data:
            return []

        columns = ", ".join(data[0].keys())
        returning_clause = f" RETURNING {', '.join(returning)}" if returning else ""

        # Build a single multi-VALUES INSERT so RETURNING captures all rows
        value_rows: List[str] = []
        all_params: Dict[str, Any] = {}
        for i, row in enumerate(data):
            safe_row = _serialize_pg_params(row)
            placeholders = ", ".join([f":{key}_{i}" for key in safe_row.keys()])
            value_rows.append(f"({placeholders})")
            for key, value in safe_row.items():
                all_params[f"{key}_{i}"] = value

        sql = (
            f"INSERT INTO {table} ({columns}) "
            f"VALUES {', '.join(value_rows)}"
            f"{returning_clause}"
        )

        async with self.session_maker() as session:
            result = await session.execute(text(sql), all_params)
            await session.commit()
            return result.mappings().all() if returning else []

    async def _update(
        self,
        table: str,
        conditions: Dict[str, Any],
        values: Dict[str, Any],
        limit: Optional[int] = None
    ) -> int:
        if not values:
            return 0

        safe_values = _serialize_pg_params(values)
        set_clause = ", ".join([f"{key} = :{key}" for key in safe_values.keys()])
        where_clause, params = self._format_conditions(conditions)
        sql = f"UPDATE {table} SET {set_clause}"

        if limit == 1:
            selection = f"SELECT ctid FROM {table}"
            if where_clause:
                selection += f" WHERE {where_clause}"
            sql = (
                f"WITH selected_rows AS ({selection} LIMIT 1) "
                f"UPDATE {table} SET {set_clause} WHERE ctid IN (SELECT ctid FROM selected_rows)"
            )

        if where_clause and limit != 1:
            sql += f" WHERE {where_clause}"

        async with self.session_maker() as session:
            result = await session.execute(text(sql), {**params, **safe_values})
            await session.commit()
            return result.rowcount

    async def _delete(
        self,
        table: str,
        conditions: Dict[str, Any],
        limit: Optional[int] = None
    ) -> int:
        where_clause, params = self._format_conditions(conditions)
        sql = f"DELETE FROM {table}"

        if limit == 1:
            selection = f"SELECT ctid FROM {table}"
            if where_clause:
                selection += f" WHERE {where_clause}"
            sql = (
                f"WITH selected_rows AS ({selection} LIMIT 1) "
                f"DELETE FROM {table} WHERE ctid IN (SELECT ctid FROM selected_rows)"
            )
        elif where_clause:
            sql += f" WHERE {where_clause}"

        async with self.session_maker() as session:
            result = await session.execute(text(sql), params)
            await session.commit()
            return result.rowcount

    def _build_select_statement(self, query: Query, for_explain: bool = False) -> Tuple[Select, Dict]:
        # Build SELECT with expression support
        select_cols = [text(col) for col in query.selects]
        stmt: Select = select(*select_cols).select_from(text(query.table))

        # -- JOINs ----------------------------------------------------------
        for join in query.joins:
            join_kw = {}
            if join.join_type.upper() == "LEFT":
                join_kw["isouter"] = True
            elif join.join_type.upper() == "FULL":
                join_kw["full"] = True
            stmt = stmt.join(
                text(join.table),
                text(join.on_clause),
                **join_kw,
            )

        where_clause = self._build_where_clause(query.wheres)
        if where_clause is not None:
            stmt = stmt.where(where_clause)

        # -- GROUP BY -------------------------------------------------------
        for col in query.group_by:
            stmt = stmt.group_by(text(col))

        for col in query.order_by:
            if col.startswith("-"):
                stmt = stmt.order_by(text(f"{col[1:]} DESC"))
            else:
                stmt = stmt.order_by(text(col))

        if query.limit is not None:
            stmt = stmt.limit(query.limit)
        if query.offset is not None:
            stmt = stmt.offset(query.offset)

        params = dict(query._params) if hasattr(query, '_params') else {}
        return stmt, params

    def _build_where_clause(self, conditions: List[WhereCondition]):
        if not conditions:
            return None
        
        and_clauses = []
        or_clauses = []
        
        for cond in conditions:
            if cond.group is not None:
                group_clause = self._build_group_clause(cond.group)
                if group_clause is not None:
                    if cond.is_or:
                        or_clauses.append(group_clause)
                    else:
                        and_clauses.append(group_clause)
            else:
                clause = self._build_single_condition(cond)
                if clause is not None:
                    if cond.is_or:
                        or_clauses.append(clause)
                    else:
                        and_clauses.append(clause)
        
        # Combine AND and OR clauses
        result_clauses = []
        if and_clauses:
            result_clauses.append(and_(*and_clauses) if len(and_clauses) > 1 else and_clauses[0])
        if or_clauses:
            result_clauses.append(or_(*or_clauses) if len(or_clauses) > 1 else or_clauses[0])
        
        if not result_clauses:
            return None
        if len(result_clauses) == 1:
            return result_clauses[0]
        return or_(*result_clauses)
    
    def _build_group_clause(self, group: WhereGroup):
        """Build a WHERE clause for a group of conditions"""
        if not group.conditions:
            return None
        
        and_clauses = []
        or_clauses = []
        for cond in group.conditions:
            if cond.group is not None:
                clause = self._build_group_clause(cond.group)
            else:
                clause = self._build_single_condition(cond)
            
            if clause is not None:
                if cond.is_or:
                    or_clauses.append(clause)
                else:
                    and_clauses.append(clause)
        
        result_clauses = []
        if and_clauses:
            result_clauses.append(and_(*and_clauses) if len(and_clauses) > 1 else and_clauses[0])
        if or_clauses:
            result_clauses.append(or_(*or_clauses) if len(or_clauses) > 1 else or_clauses[0])
        
        if not result_clauses:
            return None
        if len(result_clauses) == 1:
            return result_clauses[0]
        return or_(*result_clauses)

    def _build_single_condition(self, cond: WhereCondition):
        if cond.column is None:
            return None
        col = literal_column(cond.column)
        op = cond.operator

        if isinstance(op, Eq):    return col == self._safe_value(op.value)
        elif isinstance(op, Gt):  return col > self._safe_value(op.value)
        elif isinstance(op, Gte): return col >= self._safe_value(op.value)
        elif isinstance(op, Lt):  return col < self._safe_value(op.value)
        elif isinstance(op, Lte): return col <= self._safe_value(op.value)
        elif isinstance(op, NotEq): return col != self._safe_value(op.value)
        elif isinstance(op, In):  return col.in_(op.values)
        elif isinstance(op, Like): return col.like(op.pattern)
        elif isinstance(op, ILike): return col.ilike(op.pattern)
        elif isinstance(op, IsNull):     return col.is_(None)
        elif isinstance(op, IsNotNull):  return col.isnot(None)
        elif isinstance(op, RawExpr):
            return text(f"{cond.column} {op.sql}")
        return col == getattr(op, 'value', op)

    @staticmethod
    def _safe_value(value: Any) -> Any:
        """Convert UUIDs to strings so SQLAlchemy can render them as literals."""
        import uuid as _uuid
        if isinstance(value, _uuid.UUID):
            return str(value)
        return value
    
    async def _bulk_update(self, table: str, ids: List[Any], values: Dict[str, Any]) -> int:
        """Update multiple rows by IDs"""
        safe_values = _serialize_pg_params(values)
        set_clause = ", ".join([f"{k} = :{k}" for k in safe_values.keys()])
        sql = f"UPDATE {table} SET {set_clause} WHERE id = ANY(:ids)"
        params = {**safe_values, "ids": ids}
        
        async with self.session_maker() as session:
            result = await session.execute(text(sql), params)
            await session.commit()
            return result.rowcount
    
    async def _bulk_delete(self, table: str, ids: List[Any]) -> int:
        """Delete multiple rows by IDs"""
        sql = f"DELETE FROM {table} WHERE id = ANY(:ids)"
        
        async with self.session_maker() as session:
            result = await session.execute(text(sql), {"ids": ids})
            await session.commit()
            return result.rowcount


async def ensure_pg_database_exists(url: str) -> None:
    """Ensure the target PostgreSQL database exists, creating it if necessary.

    Parses the database name from *url*, connects to the ``postgres`` maintenance
    database, checks ``pg_database``, and issues ``CREATE DATABASE`` if the target
    is missing.  Uses ``AUTOCOMMIT`` isolation because ``CREATE DATABASE`` cannot
    run inside a transaction block.

    Args:
        url: A full ``postgresql+asyncpg://user:pass@host:port/dbname`` URL.
    """
    parsed = urlparse(url)
    db_name = parsed.path.lstrip("/")
    if not db_name:
        logger.warning("No database name in URL %r; skipping ensure-database check.", url)
        return

    # Build a maintenance URL pointing at the always-present 'postgres' database
    maintenance_url = urlunparse(parsed._replace(path="/postgres"))

    engine = create_async_engine(maintenance_url, pool_pre_ping=True)
    try:
        async with engine.connect() as conn:
            # Check if target database already exists
            result = await conn.execute(
                text("SELECT 1 FROM pg_database WHERE datname = :name"),
                {"name": db_name},
            )
            if result.fetchone() is not None:
                logger.debug("Database %r already exists.", db_name)
                return

            # CREATE DATABASE requires AUTOCOMMIT isolation
            conn_with_ac = conn.execution_options(isolation_level="AUTOCOMMIT")
            await conn_with_ac.execute(text(f"CREATE DATABASE {_pg_quote_ident(db_name)}"))
            logger.info("Created database %r.", db_name)
    finally:
        await engine.dispose()


def _pg_quote_ident(name: str) -> str:
    """Quote a PostgreSQL identifier with double-quotes, escaping internal quotes."""
    return '"' + name.replace('"', '""') + '"'
