"""
Database loaders for ETL pipelines.
"""

from __future__ import annotations

import logging
import time
from typing import Any

from sqlalchemy import create_engine, text
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
from sqlalchemy.orm import sessionmaker

from pycharter.etl_generator.database import (
    DB_MSSQL,
    DB_MYSQL,
    DB_POSTGRESQL,
    DB_SQLITE,
    DEFAULT_TUNNEL_LOCAL_PORT,
    create_ssh_tunnel,
    detect_database_type,
    load_data_mssql,
    load_data_mysql,
    load_data_postgresql,
    load_data_sqlite,
    modify_url_for_tunnel,
)
from pycharter.etl_generator.loaders.base import BaseLoader
from pycharter.etl_generator.result import LoadResult

logger = logging.getLogger(__name__)


def _ssh_tunnel_enabled(ssh_tunnel: dict[str, Any]) -> bool:
    """Return True only if ssh_tunnel is configured and enabled is truthy.
    Value injection may set enabled to the string 'true' or 'false'; treat only true-ish as enabled.
    """
    if not ssh_tunnel:
        return False
    enabled = ssh_tunnel.get("enabled")
    if enabled is None:
        return False
    return enabled in (True, "true", "1", "yes")


class PostgresLoader(BaseLoader):
    """
    Loader for PostgreSQL databases.

    Supports:
    - Insert, upsert, replace, update, delete, truncate_and_load
    - Bulk operations for efficiency
    - SSH tunneling

    Example:
        >>> loader = PostgresLoader(
        ...     connection_string="postgresql://user:pass@localhost/db",
        ...     table="users",
        ...     write_method="upsert",
        ...     primary_key="id",
        ... )
        >>> result = await loader.load(data)
    """

    def __init__(
        self,
        connection_string: str,
        table: str,
        schema: str = "public",
        write_method: str = "upsert",
        primary_key: str | list[str] | None = None,
        batch_size: int = 1000,
        ssh_tunnel: dict[str, Any] | None = None,
        update_columns: list[str] | None = None,
    ):
        self.connection_string = connection_string
        self.table = table
        self.schema = schema
        self.write_method = write_method
        self.primary_key = primary_key
        self.batch_size = batch_size
        self.ssh_tunnel = ssh_tunnel
        self.update_columns = update_columns

    @classmethod
    def from_config(cls, config: dict[str, Any]) -> "PostgresLoader":
        """Create loader from configuration dict."""
        db_config = config.get("database", {})
        return cls(
            connection_string=db_config.get("url"),
            table=config.get("table"),
            schema=config.get("schema_name", "public"),
            write_method=config.get("write_method", "upsert"),
            primary_key=config.get("primary_key"),
            batch_size=config.get("batch_size", 1000),
            ssh_tunnel=config.get("ssh_tunnel"),
            update_columns=config.get("update_columns"),
        )

    async def load(self, data: list[dict[str, Any]], **params) -> LoadResult:
        """Load data to PostgreSQL."""
        start_time = time.time()

        if not data:
            return LoadResult(success=True, rows_loaded=0)

        # Handle SSH tunnel if configured
        tunnel = None
        connection_string = self.connection_string

        tunnel_enabled = _ssh_tunnel_enabled(self.ssh_tunnel)
        if tunnel_enabled:
            tunnel = create_ssh_tunnel(self.ssh_tunnel)
            if tunnel:
                local_port = int(
                    self.ssh_tunnel.get("local_port", DEFAULT_TUNNEL_LOCAL_PORT)
                )
                connection_string = modify_url_for_tunnel(
                    connection_string, local_port, DB_POSTGRESQL
                )

        try:
            # Use async engine for PostgreSQL
            # Convert sync URL to async if needed
            if "+asyncpg" not in connection_string:
                async_url = connection_string.replace(
                    "postgresql://", "postgresql+asyncpg://"
                )
            else:
                async_url = connection_string

            # Longer connect timeout for tunnel/slow networks; asyncpg accepts timeout in seconds
            connect_args = {"timeout": 60}
            engine = create_async_engine(
                async_url, echo=False, connect_args=connect_args
            )
            async_session = sessionmaker(
                engine, class_=AsyncSession, expire_on_commit=False
            )

            async with async_session() as session:
                result = await load_data_postgresql(
                    data=data,
                    session=session,
                    schema_name=self.schema,
                    table_name=self.table,
                    write_method=self.write_method,
                    primary_key=self.primary_key,
                    batch_size=self.batch_size,
                    update_columns=getattr(self, "update_columns", None),
                )

            await engine.dispose()

            duration = time.time() - start_time
            logger.debug(
                "Loaded %s records to %s.%s in %.2fs",
                result["total"],
                self.schema,
                self.table,
                duration,
            )

            inserted = result.get("inserted", 0)
            updated = result.get("updated", 0)
            return LoadResult(
                success=True,
                rows_loaded=inserted + updated,
                inserted=inserted,
                updated=updated,
                duration_seconds=duration,
            )

        except Exception as e:
            # Log connection target (redact password) for timeout/connection errors
            try:
                from urllib.parse import urlparse

                parsed = urlparse(connection_string)
                host_port = f"{parsed.hostname or '?'}:{parsed.port or '?'}"
            except Exception:
                host_port = "?"
            err_msg = str(e)
            if len(err_msg) > 500:
                err_msg = err_msg[:500] + " ... [truncated]"
            logger.error(
                "PostgreSQL load failed: %s (connection target %s)",
                err_msg,
                host_port,
                exc_info=False,
            )
            return LoadResult(
                success=False,
                error=err_msg,
                duration_seconds=time.time() - start_time,
            )
        finally:
            if tunnel:
                tunnel.stop()


class DatabaseLoader(BaseLoader):
    """
    Generic database loader that auto-detects database type.

    Supports PostgreSQL, MySQL, SQLite, and MSSQL.

    Example:
        >>> loader = DatabaseLoader(
        ...     connection_string="mysql://user:pass@localhost/db",
        ...     table="users",
        ... )
        >>> result = await loader.load(data)
    """

    def __init__(
        self,
        connection_string: str,
        table: str,
        schema: str | None = None,
        write_method: str = "upsert",
        primary_key: str | list[str] | None = None,
        batch_size: int = 1000,
        ssh_tunnel: dict[str, Any] | None = None,
        update_columns: list[str] | None = None,
    ):
        self.connection_string = connection_string
        self.table = table
        self.schema = schema
        self.write_method = write_method
        self.primary_key = primary_key
        self.batch_size = batch_size
        self.ssh_tunnel = ssh_tunnel
        self.update_columns = update_columns
        self.db_type = detect_database_type(connection_string)

    @classmethod
    def from_config(cls, config: dict[str, Any]) -> "DatabaseLoader":
        """Create loader from configuration dict."""
        db_config = config.get("database", {})
        return cls(
            connection_string=db_config.get("url"),
            table=config.get("table"),
            schema=config.get("schema_name", "public"),
            write_method=config.get("write_method", "upsert"),
            primary_key=config.get("primary_key"),
            batch_size=config.get("batch_size", 1000),
            ssh_tunnel=config.get("ssh_tunnel"),
            update_columns=config.get("update_columns"),
        )

    async def load(self, data: list[dict[str, Any]], **params) -> LoadResult:
        """Load data using appropriate database loader."""
        if self.db_type == DB_POSTGRESQL:
            loader = PostgresLoader(
                connection_string=self.connection_string,
                table=self.table,
                schema=self.schema or "public",
                write_method=self.write_method,
                primary_key=self.primary_key,
                batch_size=self.batch_size,
                ssh_tunnel=self.ssh_tunnel,
                update_columns=getattr(self, "update_columns", None),
            )
            return await loader.load(data, **params)
        else:
            # For non-PostgreSQL databases, use sync loading
            return await self._load_sync(data, **params)

    async def _load_sync(self, data: list[dict[str, Any]], **params) -> LoadResult:
        """Load data using sync database operations."""
        start_time = time.time()

        if not data:
            return LoadResult(success=True, rows_loaded=0)

        # Handle SSH tunnel if configured
        tunnel = None
        connection_string = self.connection_string

        if _ssh_tunnel_enabled(self.ssh_tunnel):
            tunnel = create_ssh_tunnel(self.ssh_tunnel)
            if tunnel:
                local_port = int(
                    self.ssh_tunnel.get("local_port", DEFAULT_TUNNEL_LOCAL_PORT)
                )
                connection_string = modify_url_for_tunnel(
                    connection_string, local_port, self.db_type
                )

        try:
            engine = create_engine(connection_string, echo=False)
            Session = sessionmaker(bind=engine)
            session = Session()

            # Select appropriate load function
            if self.db_type == DB_MYSQL:
                result = load_data_mysql(
                    data,
                    session,
                    self.schema or "",
                    self.table,
                    self.write_method,
                    self.primary_key,
                    self.batch_size,
                )
            elif self.db_type == DB_SQLITE:
                result = load_data_sqlite(
                    data,
                    session,
                    "",
                    self.table,
                    self.write_method,
                    self.primary_key,
                    self.batch_size,
                )
            elif self.db_type == DB_MSSQL:
                result = load_data_mssql(
                    data,
                    session,
                    self.schema or "dbo",
                    self.table,
                    self.write_method,
                    self.primary_key,
                    self.batch_size,
                )
            else:
                raise ValueError(f"Unsupported database type: {self.db_type}")

            session.close()
            engine.dispose()

            duration = time.time() - start_time
            inserted = result.get("inserted", 0)
            updated = result.get("updated", 0)
            return LoadResult(
                success=True,
                rows_loaded=inserted + updated,
                inserted=inserted,
                updated=updated,
                duration_seconds=duration,
            )

        except Exception as e:
            err_msg = str(e)
            if len(err_msg) > 500:
                err_msg = err_msg[:500] + " ... [truncated]"
            logger.error("Database load failed: %s", err_msg, exc_info=False)
            return LoadResult(
                success=False,
                error=err_msg,
                duration_seconds=time.time() - start_time,
            )
        finally:
            if tunnel:
                tunnel.stop()
