"""Local PostgreSQL operations — create tables, insert data, query v3q_data schema."""

from __future__ import annotations

import logging
import math
import time
from datetime import date, datetime
from typing import Any

import numpy as np
import pandas as pd

from database.connection_manager import ConnectionManager

logger = logging.getLogger(__name__)


def _sanitize_value(v: Any) -> Any:
    """Convert pandas/numpy types to native Python types for asyncpg."""
    if v is None:
        return None
    if isinstance(v, float) and math.isnan(v):
        return None
    if isinstance(v, type(pd.NaT)):
        return None
    if isinstance(v, (np.bool_,)):
        return bool(v)
    if isinstance(v, (np.integer,)):
        return int(v)
    if isinstance(v, (np.floating,)):
        f = float(v)
        return None if math.isnan(f) else f
    if isinstance(v, (np.datetime64,)):
        if pd.isna(v):
            return None
        return pd.Timestamp(v).to_pydatetime()
    if isinstance(v, pd.Timestamp):
        if pd.isna(v):
            return None
        return v.to_pydatetime()
    return v


def _sanitize_record(record: dict) -> dict:
    """Sanitize all values in a record dict for asyncpg insertion."""
    return {k: _sanitize_value(v) for k, v in record.items()}


class LocalPG:
    """High-level operations on the local PostgreSQL v3q_data schema."""

    def __init__(self, cm: ConnectionManager) -> None:
        self._cm = cm

    async def create_table_from_columns(
        self,
        table_name: str,
        columns: list[dict[str, str]],
        schema: str = "v3q_data",
    ) -> None:
        """Create a table in v3q_data from column definitions.

        columns: [{"name": "col1", "type": "TEXT"}, ...]
        """
        col_defs = ", ".join(
            f'"{col["name"]}" {col["type"]}' for col in columns
        )
        ddl = f'CREATE TABLE IF NOT EXISTS {schema}."{table_name}" ({col_defs})'
        logger.info(f"[LOCAL] CREATE TABLE {schema}.{table_name} ({len(columns)} columns)")
        t0 = time.perf_counter()
        await self._cm.execute_write(ddl)
        logger.info(f"[LOCAL] Table created in {time.perf_counter() - t0:.2f}s")

    async def drop_table(self, table_name: str, schema: str = "v3q_data") -> None:
        logger.info(f"[LOCAL] DROP TABLE {schema}.{table_name}")
        await self._cm.execute_write(f'DROP TABLE IF EXISTS {schema}."{table_name}" CASCADE')

    async def truncate_table(self, table_name: str, schema: str = "v3q_data") -> None:
        logger.info(f"[LOCAL] TRUNCATE TABLE {schema}.{table_name}")
        await self._cm.execute_write(f'TRUNCATE TABLE {schema}."{table_name}"')

    async def insert_dataframe(
        self,
        table_name: str,
        df: pd.DataFrame,
        schema: str = "v3q_data",
        batch_size: int = 10000,
    ) -> int:
        """Insert a pandas DataFrame into a local PG table.

        Uses asyncpg COPY protocol (binary bulk load) for performance.
        Falls back to positional-parameter INSERT if COPY fails.
        Handles SAP column names with special characters (/, etc.).
        """
        if df.empty:
            logger.debug(f"[LOCAL] insert_dataframe: empty DataFrame for {schema}.{table_name}, skipping")
            return 0

        mem_mb = df.memory_usage(deep=True).sum() / (1024 * 1024)
        logger.info(
            f"[LOCAL] insert_dataframe: {len(df):,} rows x {len(df.columns)} cols → {schema}.{table_name} "
            f"(batch_size={batch_size}, df_memory={mem_mb:.1f}MB)"
        )
        t0 = time.perf_counter()

        try:
            count = await self._insert_copy(table_name, df, schema, batch_size)
            elapsed = time.perf_counter() - t0
            rate = count / elapsed if elapsed > 0 else 0
            logger.info(
                f"[LOCAL] COPY insert complete: {count:,} rows in {elapsed:.2f}s ({rate:,.0f} rows/s)"
            )
            return count
        except Exception as e:
            elapsed = time.perf_counter() - t0
            logger.warning(
                f"[LOCAL] COPY insert FAILED after {elapsed:.2f}s ({type(e).__name__}: {e}), "
                f"falling back to positional INSERT"
            )
            t1 = time.perf_counter()
            count = await self._insert_positional(table_name, df, schema, batch_size)
            elapsed2 = time.perf_counter() - t1
            rate = count / elapsed2 if elapsed2 > 0 else 0
            logger.info(
                f"[LOCAL] Positional INSERT complete: {count:,} rows in {elapsed2:.2f}s ({rate:,.0f} rows/s)"
            )
            return count

    async def _insert_copy(
        self, table_name: str, df: pd.DataFrame, schema: str, batch_size: int
    ) -> int:
        """Bulk insert via asyncpg COPY protocol — 10-100x faster than row-by-row INSERT."""
        columns = list(df.columns)
        total_rows = 0
        batch_num = 0
        num_batches = (len(df) + batch_size - 1) // batch_size
        logger.debug(f"[LOCAL] COPY: {len(df):,} rows in {num_batches} batches → {schema}.{table_name}")

        async with self._cm.raw_asyncpg() as raw_conn:
            for start in range(0, len(df), batch_size):
                batch_num += 1
                t_batch = time.perf_counter()
                batch = df.iloc[start : start + batch_size]

                t_sanitize = time.perf_counter()
                records = [
                    tuple(_sanitize_value(v) for v in row)
                    for row in batch.itertuples(index=False, name=None)
                ]
                sanitize_elapsed = time.perf_counter() - t_sanitize

                t_copy = time.perf_counter()
                await raw_conn.copy_records_to_table(
                    table_name,
                    records=records,
                    columns=columns,
                    schema_name=schema,
                )
                copy_elapsed = time.perf_counter() - t_copy

                total_rows += len(records)
                batch_elapsed = time.perf_counter() - t_batch
                logger.info(
                    f"[LOCAL] COPY batch {batch_num}/{num_batches}: {len(records):,} rows in {batch_elapsed:.2f}s "
                    f"(sanitize={sanitize_elapsed:.2f}s, copy={copy_elapsed:.2f}s, total={total_rows:,})"
                )

        return total_rows

    async def _insert_positional(
        self, table_name: str, df: pd.DataFrame, schema: str, batch_size: int
    ) -> int:
        """Fallback INSERT using $1,$2,... positional params (handles special column names)."""
        columns = [f'"{c}"' for c in df.columns]
        col_list = ", ".join(columns)
        placeholders = ", ".join(f"${i + 1}" for i in range(len(df.columns)))
        query = f'INSERT INTO {schema}."{table_name}" ({col_list}) VALUES ({placeholders})'
        logger.debug(f"[LOCAL] Positional INSERT query: {query[:120]}...")

        total_rows = 0
        batch_num = 0
        num_batches = (len(df) + batch_size - 1) // batch_size

        async with self._cm.raw_asyncpg() as raw_conn:
            for start in range(0, len(df), batch_size):
                batch_num += 1
                t_batch = time.perf_counter()
                batch = df.iloc[start : start + batch_size]
                records = [
                    tuple(_sanitize_value(v) for v in row)
                    for row in batch.itertuples(index=False, name=None)
                ]
                await raw_conn.executemany(query, records)
                total_rows += len(records)
                batch_elapsed = time.perf_counter() - t_batch
                logger.info(
                    f"[LOCAL] INSERT batch {batch_num}/{num_batches}: {len(records):,} rows in {batch_elapsed:.2f}s "
                    f"(total={total_rows:,})"
                )

        return total_rows

    async def execute_readonly(self, sql: str) -> tuple[list[dict], list[str]]:
        """Execute a read-only SQL query against local PG.

        Returns (rows_as_dicts, column_names).
        The query is wrapped in a read-only transaction.
        """
        wrapped = f"SET TRANSACTION READ ONLY; {sql}"
        async with self._cm.session() as session:
            # Set read-only
            await session.execute(
                __import__("sqlalchemy").text("SET TRANSACTION READ ONLY")
            )
            from sqlalchemy import text
            result = await session.execute(text(sql))
            columns = list(result.keys())
            rows = [dict(zip(columns, row)) for row in result.fetchall()]
            return rows, columns

    async def get_row_count(self, table_name: str, schema: str = "v3q_data") -> int:
        row = await self._cm.fetch_one(
            f'SELECT COUNT(*) as cnt FROM {schema}."{table_name}"'
        )
        return row["cnt"] if row else 0
