"""FHIR SQL client helpers for Postgres-backed FHIR databases."""

from __future__ import annotations

import logging
import os
import multiprocessing as mp
import re
from dataclasses import dataclass
from typing import Any, Dict, Iterable, Iterator, List, Optional, Sequence, Tuple, Union

import pandas as pd
try:
    import psycopg
except Exception as exc:  # pragma: no cover
    psycopg = None  # type: ignore
    _PSYCOPG_IMPORT_ERROR = exc

logger = logging.getLogger(__name__)


class SQLQueryError(RuntimeError):
    """Raised when a SQL query fails or returns invalid results."""


@dataclass
class SQLConfig:
    dsn: Optional[str]
    host: Optional[str]
    port: Optional[int]
    dbname: Optional[str]
    user: Optional[str]
    password: Optional[str]
    sslmode: Optional[str]


class FHIRSQLClient:
    """Lightweight SQL client for FHIR-backed databases.

    Environment variables (optional):
    - FHIR_DB_DSN
    - FHIR_DB_HOST
    - FHIR_DB_PORT
    - FHIR_DB_NAME
    - FHIR_DB_USER
    - FHIR_DB_PASSWORD
    - FHIR_DB_SSLMODE
    """

    def __init__(
        self,
        dsn: Optional[str] = None,
        host: Optional[str] = None,
        port: Optional[int] = None,
        dbname: Optional[str] = None,
        user: Optional[str] = None,
        password: Optional[str] = None,
        sslmode: Optional[str] = None,
        connect_timeout: int = 10,
        log_queries: bool = False,
    ) -> None:
        """Initialize a SQL client for Postgres-backed FHIR data."""
        if psycopg is None:
            raise ImportError(
                "psycopg is required for FHIRSQLClient. Install with `pip install psycopg` "
                "and ensure libpq is available."
            ) from _PSYCOPG_IMPORT_ERROR
        self.config = SQLConfig(
            dsn=dsn or os.environ.get("FHIR_DB_DSN"),
            host=host or os.environ.get("FHIR_DB_HOST"),
            port=port or _maybe_int(os.environ.get("FHIR_DB_PORT")),
            dbname=dbname or os.environ.get("FHIR_DB_NAME"),
            user=user or os.environ.get("FHIR_DB_USER"),
            password=password or os.environ.get("FHIR_DB_PASSWORD"),
            sslmode=sslmode or os.environ.get("FHIR_DB_SSLMODE"),
        )
        self.connect_timeout = connect_timeout
        self.log_queries = log_queries
        self._conn = self._connect()

    def __enter__(self) -> "FHIRSQLClient":
        """Return self for context manager usage."""
        return self

    def close(self) -> None:
        """Close the database connection."""
        if self._conn is not None:
            self._conn.close()

    def __exit__(self, exc_type, exc, exc_tb) -> None:
        """Close the connection on context manager exit."""
        self.close()

    def query(
        self,
        sql: str,
        params: Optional[Dict[str, Any] | Sequence[Any]] = None,
    ) -> List[Dict[str, Any]]:
        """Execute a SQL query and return rows as dicts."""
        try:
            with self._conn.cursor() as cur:
                if self.log_queries:
                    logger.info("SQL query execute")
                    logger.debug("SQL: %s", sql)
                cur.execute(sql, params)
                if cur.description is None:
                    return []
                columns = [col.name for col in cur.description]
                rows = cur.fetchall()
            return [dict(zip(columns, row)) for row in rows]
        except psycopg.Error as exc:
            logger.exception("SQL query failed")
            raise SQLQueryError("SQL query failed") from exc

    def query_to_dataframe(
        self,
        sql: str,
        params: Optional[Dict[str, Any] | Sequence[Any]] = None,
        rename: Optional[Dict[str, str]] = None,
    ) -> pd.DataFrame:
        """Execute a SQL query and return a DataFrame.

        Parameters
        ----------
        sql:
            SQL statement with placeholders.
        params:
            Optional parameters (dict for named, sequence for positional).
        rename:
            Optional mapping to rename columns in the resulting DataFrame.
        """
        try:
            if self.log_queries:
                logger.info("SQL query_to_dataframe execute")
                logger.debug("SQL: %s", sql)
            df = pd.read_sql_query(sql, self._conn, params=params)
            if rename:
                df = df.rename(columns=rename)
            return df
        except Exception as exc:
            logger.exception("SQL query_to_dataframe failed")
            raise SQLQueryError("SQL query_to_dataframe failed") from exc

    def query_df(
        self,
        sql: str,
        params: Optional[Dict[str, Any] | Sequence[Any]] = None,
        rename: Optional[Dict[str, str]] = None,
    ) -> pd.DataFrame:
        """Backward-compatible alias for query_to_dataframe."""
        return self.query_to_dataframe(sql=sql, params=params, rename=rename)

    def query_df_stream(
        self,
        sql: str,
        params: Optional[Dict[str, Any] | Sequence[Any]] = None,
        rename: Optional[Dict[str, str]] = None,
        chunksize: int = 10000,
    ) -> Iterator[pd.DataFrame]:
        """Stream a large SQL query as DataFrame chunks."""
        try:
            if self.log_queries:
                logger.info("SQL query_df_stream execute")
                logger.debug("SQL: %s", sql)
            for chunk in pd.read_sql_query(sql, self._conn, params=params, chunksize=chunksize):
                if rename:
                    chunk = chunk.rename(columns=rename)
                yield chunk
        except Exception as exc:
            logger.exception("SQL query_df_stream failed")
            raise SQLQueryError("SQL query_df_stream failed") from exc

    def query_df_parallel(
        self,
        jobs: Sequence[Tuple[str, Optional[Dict[str, Any] | Sequence[Any]]]],
        rename: Optional[Dict[str, str]] = None,
        max_workers: Optional[int] = None,
        merge: str = "concat",
        unique_on: Optional[Sequence[str]] = None,
    ) -> pd.DataFrame:
        """Run multiple SQL queries in parallel and merge results.

        Parameters
        ----------
        jobs:
            List of (sql, params) pairs.
        rename:
            Optional mapping to rename output columns.
        max_workers:
            Number of worker processes (defaults to CPU count).
        """
        if not jobs:
            return pd.DataFrame()
        if self.log_queries:
            logger.info("SQL query_df_parallel jobs=%s", len(jobs))
        worker_args = [
            (self.config, self.connect_timeout, sql, params, rename)
            for sql, params in jobs
        ]
        try:
            with mp.Pool(processes=max_workers) as pool:
                frames = pool.map(_worker_query_df, worker_args)
        except Exception as exc:
            logger.exception("Parallel SQL execution failed")
            raise SQLQueryError("Parallel SQL execution failed") from exc
        if not frames:
            return pd.DataFrame()
        return _merge_frames(frames, merge=merge, unique_on=unique_on)

    def query_from_dataframe(
        self,
        df: pd.DataFrame,
        sql: str,
        column_map: Dict[str, str],
        rename: Optional[Dict[str, str]] = None,
        param_limit: int = 50000,
        chunk_size: Optional[int] = None,
        parallel: bool = False,
        max_workers: Optional[int] = None,
        merge: str = "concat",
        unique_on: Optional[Sequence[str]] = None,
        param_style: str = "positional",
        param_safety_factor: float = 0.9,
    ) -> pd.DataFrame:
        """Run a query with parameters sourced from a DataFrame.

        The SQL must contain named placeholder blocks using `{name}` or `{{name}}` for each column_map key.
        Example:
        `SELECT * FROM observation WHERE patient_id IN ({patient_ids})`

        Parameters
        ----------
        df:
            Input DataFrame containing values.
        sql:
            SQL statement with `{name}` placeholder blocks.
        column_map:
            Mapping from placeholder name -> DataFrame column name.
        rename:
            Optional mapping to rename output columns.
        param_limit:
            Max number of parameters per query (Postgres limit is 65535).
        chunk_size:
            Optional explicit chunk size for splitting the DataFrame.
        parallel:
            Whether to execute chunks in parallel and merge results.
        max_workers:
            Number of worker processes when parallel is True.
        merge:
            How to merge results: "concat" or "drop_duplicates".
        unique_on:
            Columns to use for duplicate dropping when merge="drop_duplicates".
        param_style:
            "positional" (default) or "named".
        param_safety_factor:
            Fraction of param_limit used for chunk sizing (0.0-1.0).
        """
        if self.log_queries:
            logger.info(
                "SQL query_from_dataframe rows=%s chunk_size=%s parallel=%s",
                len(df),
                chunk_size,
                parallel,
            )
        if df.empty:
            return pd.DataFrame()
        for placeholder, col in column_map.items():
            if col not in df.columns:
                raise ValueError(f"Column '{col}' not found in DataFrame.")
            if df[col].isnull().any():
                raise ValueError(f"Column '{col}' contains null values.")
            if f"{{{placeholder}}}" not in sql and f"{{{{{placeholder}}}}}" not in sql:
                raise ValueError(
                    f"Missing placeholder '{{{placeholder}}}' or '{{{{{placeholder}}}}}' in SQL."
                )
        if param_style not in {"positional", "named"}:
            raise ValueError("param_style must be 'positional' or 'named'.")
        if not (0.0 < param_safety_factor <= 1.0):
            raise ValueError("param_safety_factor must be within (0.0, 1.0].")

        col_count = max(1, len(column_map))
        if chunk_size is None:
            chunk_size = max(1, int((param_limit * param_safety_factor) // col_count))

        frames: List[pd.DataFrame] = []
        if parallel:
            jobs: List[Tuple[str, Union[Sequence[Any], Dict[str, Any]]]] = []
            for chunk in _iter_chunks(df, chunk_size):
                query_sql, params = _expand_sql_placeholders(
                    sql=sql,
                    df=chunk,
                    column_map=column_map,
                    param_style=param_style,
                )
                jobs.append((query_sql, params))
            frames = [
                self.query_df_parallel(
                    jobs=jobs,
                    rename=rename,
                    max_workers=max_workers,
                    merge=merge,
                    unique_on=unique_on,
                )
            ]
        else:
            for chunk in _iter_chunks(df, chunk_size):
                query_sql, params = _expand_sql_placeholders(
                    sql=sql,
                    df=chunk,
                    column_map=column_map,
                    param_style=param_style,
                )
                frame = self.query_to_dataframe(query_sql, params=params, rename=rename)
                frames.append(frame)
        if not frames:
            return pd.DataFrame()
        return _merge_frames(frames, merge=merge, unique_on=unique_on)

    def query_df_from(
        self,
        df: pd.DataFrame,
        sql: str,
        column_map: Dict[str, str],
        rename: Optional[Dict[str, str]] = None,
        param_limit: int = 50000,
        chunk_size: Optional[int] = None,
        parallel: bool = False,
        max_workers: Optional[int] = None,
        merge: str = "concat",
        unique_on: Optional[Sequence[str]] = None,
        param_style: str = "positional",
        param_safety_factor: float = 0.9,
    ) -> pd.DataFrame:
        """Short alias for query_from_dataframe."""
        return self.query_from_dataframe(
            df=df,
            sql=sql,
            column_map=column_map,
            rename=rename,
            param_limit=param_limit,
            chunk_size=chunk_size,
            parallel=parallel,
            max_workers=max_workers,
            merge=merge,
            unique_on=unique_on,
            param_style=param_style,
            param_safety_factor=param_safety_factor,
        )

    def _connect(self) -> psycopg.Connection:
        """Create a database connection from config."""
        try:
            if self.config.dsn:
                conn = psycopg.connect(self.config.dsn, connect_timeout=self.connect_timeout)
            else:
                if not all([self.config.host, self.config.dbname, self.config.user]):
                    raise ValueError("Missing DB connection info. Set FHIR_DB_* env vars or pass args.")
                conn = psycopg.connect(
                    host=self.config.host,
                    port=self.config.port or 5432,
                    dbname=self.config.dbname,
                    user=self.config.user,
                    password=self.config.password,
                    sslmode=self.config.sslmode,
                    connect_timeout=self.connect_timeout,
                )
            if self.log_queries:
                logger.info("Connected to SQL database")
            return conn
        except Exception as exc:
            logger.exception("Failed to connect to SQL database")
            raise SQLQueryError("Failed to connect to SQL database") from exc


def _iter_chunks(df: pd.DataFrame, size: int) -> Iterable[pd.DataFrame]:
    """Yield DataFrame chunks of a given size."""
    for start in range(0, len(df), size):
        yield df.iloc[start : start + size]


def _expand_sql_placeholders(
    sql: str,
    df: pd.DataFrame,
    column_map: Dict[str, str],
    param_style: str = "positional",
) -> Tuple[str, Union[List[Any], Dict[str, Any]]]:
    """Expand {placeholders} into SQL parameters and return SQL+params."""
    params_list: List[Any] = []
    params_named: Dict[str, Any] = {}
    for placeholder, col in column_map.items():
        values = df[col].tolist()
        if not values:
            values = [None]
        if param_style == "named":
            keys = [f"{placeholder}_{i}" for i in range(len(values))]
            placeholders = ",".join([f"%({k})s" for k in keys])
            params_named.update({k: v for k, v in zip(keys, values)})
        else:
            placeholders = ",".join(["%s"] * len(values))
            params_list.extend(values)
        sql = _safe_replace_placeholder(sql, placeholder, placeholders)
    return (sql, params_named) if param_style == "named" else (sql, params_list)


def _maybe_int(value: Optional[str]) -> Optional[int]:
    """Parse an optional int from string."""
    if value is None or value == "":
        return None
    try:
        return int(value)
    except ValueError:
        return None


def _worker_query_df(
    args: Tuple[SQLConfig, int, str, Optional[Dict[str, Any] | Sequence[Any]], Optional[Dict[str, str]]]
) -> pd.DataFrame:
    """Worker helper for parallel SQL queries."""
    config, connect_timeout, sql, params, rename = args
    conn = _connect_from_config(config, connect_timeout)
    try:
        df = pd.read_sql_query(sql, conn, params=params)
        if rename:
            df = df.rename(columns=rename)
        return df
    except Exception as exc:
        raise SQLQueryError("Parallel SQL query failed") from exc
    finally:
        conn.close()


def _connect_from_config(config: SQLConfig, connect_timeout: int) -> "psycopg.Connection":
    """Connect to Postgres using a SQLConfig."""
    if config.dsn:
        return psycopg.connect(config.dsn, connect_timeout=connect_timeout)  # type: ignore[union-attr]
    if not all([config.host, config.dbname, config.user]):
        raise ValueError("Missing DB connection info. Set FHIR_DB_* env vars or pass args.")
    return psycopg.connect(  # type: ignore[union-attr]
        host=config.host,
        port=config.port or 5432,
        dbname=config.dbname,
        user=config.user,
        password=config.password,
        sslmode=config.sslmode,
        connect_timeout=connect_timeout,
    )


def _safe_replace_placeholder(sql: str, name: str, replacement: str) -> str:
    """Replace {name} or {{name}} placeholders safely."""
    token_legacy = f"{{{name}}}"
    token_safe = f"{{{{{name}}}}}"
    has_safe = token_safe in sql
    has_legacy = token_legacy in sql
    if not has_safe and not has_legacy:
        raise ValueError(f"Missing placeholder '{{{name}}}' or '{{{{{name}}}}}' in SQL.")
    if has_safe:
        sql = sql.replace(token_safe, replacement)
    if has_legacy:
        # Avoid replacing accidental braces in strings by using regex for exact token.
        sql = re.sub(re.escape(token_legacy), replacement, sql)
    return sql


def _merge_frames(
    frames: Sequence[pd.DataFrame],
    merge: str = "concat",
    unique_on: Optional[Sequence[str]] = None,
) -> pd.DataFrame:
    """Merge a list of DataFrames with optional deduplication."""
    if not frames:
        return pd.DataFrame()
    df = pd.concat(frames, ignore_index=True)
    if merge == "drop_duplicates":
        return df.drop_duplicates(subset=list(unique_on) if unique_on else None).reset_index(drop=True)
    if merge != "concat":
        raise ValueError("merge must be 'concat' or 'drop_duplicates'.")
    return df
