from __future__ import annotations

import os
import re
import sys

from typing import TYPE_CHECKING, Any

import duckdb
# https://duckdb.org/docs/stable/clients/python/known_issues#numpy-import-multithreading
import numpy.core.multiarray  # noqa: F401
import pandas as pd
import param
import sqlglot

from ..config import config
from ..serializers import Serializer
from ..transforms import Filter
from ..transforms.sql import (
    SQLCount, SQLFilter, SQLLimit, SQLSelectFrom,
)
from ..util import detect_file_encoding, normalize_table_name
from .base import BaseSQLSource, Source, cached

if TYPE_CHECKING:
    import pandas as pd


class DuckDBSource(BaseSQLSource):
    """
    DuckDBSource provides a simple wrapper around the DuckDB SQL
    connector.

    To specify tables to be queried provide a list or dictionary of
    tables. A SQL expression to fetch the data from the table will
    then be generated using the `sql_expr`, e.g. if we specify a local
    table flights.db the default sql_expr `SELECT * FROM {table}` will
    expand that to `SELECT * FROM flights.db`. If you want to specify
    a full SQL expression as a table you must change the `sql_expr` to
    '{table}' ensuring no further templating is applied.

    Note that certain functionality in DuckDB requires modules to be
    loaded before making a query. These can be specified using the
    `initializers` parameter, providing the ability to define DuckDb
    statements to be run when initializing the connection.
    """

    filter_in_sql = param.Boolean(default=True, doc="""
        Whether to apply filters in SQL or in-memory.""")

    initializers = param.List(default=[], doc="""
        SQL statements to run to initialize the connection.""")

    mirrors = param.Dict(default={}, doc="""
        Mirrors the tables into the DuckDB database. The mirrors
        should define a mapping from the table names to the source of
        the mirror which may be defined as a Pipeline or a tuple of
        the source and the table.""")

    load_schema = param.Boolean(default=True, doc="Whether to load the schema")

    sql_expr = param.String(default='SELECT * FROM {table}', doc="""
        The SQL expression to execute.""")

    ephemeral = param.Boolean(default=False, doc="""
        Whether the data is ephemeral, i.e. manually inserted into the
        DuckDB table or derived from real data.""")

    read_only = param.Boolean(default=None, doc="""
        Whether to open the DuckDB database in read-only mode.""")

    tables = param.ClassSelector(class_=(list, dict), doc="""
        List or dictionary of tables.""")

    table_params = param.Dict(default={}, doc="""
        Dictionary mapping table names to lists of SQL parameters.
        Parameters are used with placeholders (?) in SQL expressions.""")

    uri = param.String(doc="The URI of the DuckDB database")

    source_type = 'duckdb'

    dialect = 'duckdb'

    _install_extension_pattern = re.compile(r"^\s*INSTALL\s+([A-Za-z_][\w]*)\s*;?\s*$", re.IGNORECASE)

    def __init__(self, **params):
        connection = params.pop('_connection', None)
        params["read_only"] = params.get('read_only', params.get("uri") not in (None, ':memory:'))
        super().__init__(**params)
        if connection:
            self._connection = connection
        else:
            self._connection = duckdb.connect(self.uri, read_only=self.read_only)
            for init in self.initializers:
                self._run_initializer(init)

        # Process tables to handle automatic file detection
        self._file_based_tables = {}
        if isinstance(self.tables, dict):
            processed_tables = {}
            sql_based_tables = {}
            # First pass: separate file paths from SQL expressions
            for table_name, table_expr in self.tables.items():
                if isinstance(table_expr, str) and self._is_file_path(table_expr):
                    table_name = re.sub(r'\W+', '_', table_name)
                    self._file_based_tables[table_name] = table_expr
                else:
                    sql_based_tables[table_name] = table_expr

            # Second pass: create views for file-based tables
            for table_name, file_path in self._file_based_tables.items():
                # Auto-detect file type and create appropriate view
                read_expr = self._create_file_read_expr(file_path)
                # Quote table name to handle special characters
                quoted_table = f'"{table_name}"' if not (table_name.startswith('"') and table_name.endswith('"')) else table_name
                view_sql = f"CREATE OR REPLACE VIEW {quoted_table} AS {read_expr}"
                cursor = self._connection.cursor()
                try:
                    cursor.execute(view_sql)
                    # Store the SQL expression for later use
                    processed_tables[table_name] = f"SELECT * FROM {quoted_table}"
                except Exception:
                    # If view creation fails, store the read expression directly
                    processed_tables[table_name] = read_expr
                finally:
                    cursor.close()

            # Third pass: process SQL-based tables (which may reference file-based tables)
            for table_name, sql_expr in sql_based_tables.items():
                table_alias = self.normalize_table(table_name)
                # Skip non-string expressions (e.g., dicts)
                # Prevent recursive view creation
                equivalent_sql_exprs = (
                    self.sql_expr.format(table=f'"{table_name}"'),
                    self.sql_expr.format(table=table_name),
                )
                if not isinstance(sql_expr, str) or sql_expr in equivalent_sql_exprs:
                    processed_tables[table_alias] = sql_expr
                    continue

                # For SQL expressions that define complete queries, create them as views
                # This includes both READ_* functions and other SELECT statements
                if sql_expr.strip().upper().startswith('SELECT'):
                    quoted_table = f'"{table_alias}"' if not (table_alias.startswith('"') and table_alias.endswith('"')) else table_alias
                    view_sql = f"CREATE OR REPLACE VIEW {quoted_table} AS {sql_expr}"
                    cursor = self._connection.cursor()
                    try:
                        cursor.execute(view_sql)
                    except Exception:
                        pass  # View creation failed, but we'll still use the original SQL
                    finally:
                        cursor.close()

                processed_tables[table_alias] = sql_expr

            self.tables = processed_tables

        for table, mirror in self.mirrors.items():
            if isinstance(mirror, pd.DataFrame):
                df = mirror
            elif isinstance(mirror, tuple):
                source, src_table = mirror
                df = source.get(src_table)
            else:
                df = mirror.data
                def update(e, table=table):
                    df = e.new
                    for col in df.select_dtypes(include=['string']).columns:
                        df[col] = df[col].astype(object)
                    self._connection.from_df(df).to_view(table)
                    self._set_cache(df, table)
                mirror.param.watch(update, 'data')
            try:
                for col in df.select_dtypes(include=['string']).columns:
                    df[col] = df[col].astype(object)
                self._connection.from_df(df).to_view(table)
            except (duckdb.CatalogException, duckdb.ParserException):
                continue

    @property
    def connection(self):
        return self._connection

    def _run_initializer(self, init: str) -> None:
        try:
            with self._connection.cursor() as cursor:
                cursor.execute(init)
        except duckdb.IOException as e:
            message = str(e)
            if "delete file" in message and "duckdb_extension" in message:
                return
            match = self._install_extension_pattern.match(init)
            if (
                sys.platform == 'win32'
                and match is not None
                and 'Could not move file: Access is denied.' in message
            ):
                # DuckDB extension installs can hit transient file locking on Windows.
                # If the extension is already present, LOAD succeeds and we can continue.
                extension = match.group(1)
                with self._connection.cursor() as cursor:
                    cursor.execute(f"LOAD {extension};")
                return
            raise

    def _is_file_path(self, table_expr: str) -> bool:
        """
        Detect if a table expression is a file path vs SQL expression.

        Returns True if it's likely a file path, False if it's SQL.
        """
        if not isinstance(table_expr, str):
            return False
        file_extensions = ('.csv', '.parquet', '.json', '.jsonl', '.ndjson', '.tsv')
        if any(table_expr.lower().endswith(ext) for ext in file_extensions):
            return True
        return False

    def _create_file_read_expr(self, file_path: str) -> str:
        """
        Create appropriate DuckDB read expression based on file extension.
        """
        file_lower = file_path.lower()

        if file_lower.endswith('.csv'):
            encoding = detect_file_encoding(file_path)
            return f"SELECT * FROM READ_CSV('{file_path}', encoding='{encoding}')"
        elif file_lower.endswith('.parquet'):
            return f"SELECT * FROM READ_PARQUET('{file_path}')"
        elif file_lower.endswith(('.json', '.jsonl', '.ndjson')):
            return f"SELECT * FROM READ_JSON('{file_path}')"
        elif file_lower.endswith('.tsv'):
            return f"SELECT * FROM READ_CSV('{file_path}', delim='\t')"
        else:
            # Default to CSV for unknown extensions
            return f"SELECT * FROM READ_CSV('{file_path}')"

    @classmethod
    def _recursive_resolve(
        cls, spec: dict[str, Any], source_type: type[Source]
    ) -> tuple[dict[str, Any], dict[str, Any]]:
        spec, refs = super()._recursive_resolve(spec, source_type)
        return spec, refs

    def _serialize_tables(self):
        tables = {}
        for t in self.get_tables():
            tdf = self.get(t)
            serializer = Serializer._get_type(config.serializer)()
            tables[t] = serializer.serialize(tdf)
        return tables

    def _process_sql_paths(self, sql_expr: dict | str) -> str:
        if isinstance(sql_expr, dict):
            return {self._process_sql_paths(k): v for k, v in sql_expr.items()}

        # Check if this is a raw file path (not a SQL expression)
        if self._is_file_path(sql_expr):
            if not re.match(r'^(?:http|ftp)s?://', sql_expr) and not os.path.isabs(sql_expr):
                return os.path.abspath(sql_expr)
            return sql_expr

        # Look for read_* patterns (case insensitive) like read_parquet, READ_CSV etc.
        matches = re.finditer(r"(?i)read_\w+\('([^']+)'\)", sql_expr)
        processed_sql = sql_expr

        for match in matches:
            file_path = match.group(1)
            # Ensure it isn't a URL and convert to absolute path if it's not already
            if re.match(r'^(?:http|ftp)s?://', file_path):
                continue
            if not os.path.isabs(file_path):
                abs_path = os.path.abspath(file_path)
                # Replace the old path with the absolute path in the SQL
                processed_sql = processed_sql.replace(
                    f"'{file_path}'", f"'{abs_path}'"
                )
        return processed_sql

    def to_spec(self, context: dict[str, Any] | None = None) -> dict[str, Any]:
        spec = super().to_spec(context)
        if self.ephemeral:
            spec['tables'] = self._serialize_tables()
        # Handle tables that are either a list or dictionary
        elif isinstance(self.tables, list):
            spec['tables'] = [self._process_sql_paths(table_name) for table_name in self.tables]
        else:
            # For dictionary case, process each SQL expression
            processed_tables = {}
            for table_name, sql_expr in self.tables.items():
                if table_name in self._file_based_tables:
                    sql_expr = self._file_based_tables[table_name]
                processed_tables[table_name] = self._process_sql_paths(sql_expr)
            spec['tables'] = processed_tables

        if 'mirrors' not in spec:
            return spec

        from ..pipeline import Pipeline
        mirrors = {}
        for table, mirror in spec['mirrors'].items():
            if isinstance(mirror, pd.DataFrame):
                serializer = Serializer._get_type(config.serializer)()
                mirrors[table] = serializer.serialize(mirror)
            elif isinstance(mirror, tuple):
                source, src_table = mirror
                src_spec = source.to_spec(context=context)
                mirrors[table] = (src_spec, src_table)
            elif isinstance(mirror, Pipeline):
                mirrors[table] = mirror.to_spec(context=context)
        spec['mirrors'] = mirrors
        return spec

    @classmethod
    def from_df(cls, tables: dict[str, pd.DataFrame], **kwargs):
        """
        Creates an ephemeral, in-memory DuckDBSource containing the
        supplied dataframe.

        Arguments
        ---------
        tables: dict[str, pandas.DataFrame]
            A dictionary mapping from table names to DataFrames
        kwargs: any
            Additional keyword arguments for the source

        Returns
        -------
        source: DuckDBSource
        """
        source = DuckDBSource(uri=':memory:', ephemeral=True, **kwargs)
        table_defs = {}
        for name, df in tables.items():
            # DuckDB occasionally has issues with the new pandas string dtype
            for col in df.select_dtypes(include=['string']).columns:
                df[col] = df[col].astype(object)
            source._connection.from_df(df).to_view(name)
            table_defs[name] = source.sql_expr.format(table=name)
        source.tables = table_defs
        return source

    @classmethod
    def from_spec(cls, spec: dict[str, Any] | str) -> Source:
        if spec.get('ephemeral') and 'tables' in spec:
            ephemeral_tables, tables = {}, {}
            for table, table_spec in spec['tables'].items():
                if isinstance(table_spec, str):
                    tables[table] = table_spec
                else:
                    ephemeral_tables[table] = table_spec
            spec['tables'] = tables
        else:
            ephemeral_tables = {}

        mirrors = spec.pop("mirrors", {})

        source = super().from_spec(spec)
        if not ephemeral_tables:
            return source
        new_tables = {}
        for t, data in ephemeral_tables.items():
            df = Serializer.deserialize(data)
            for col in df.select_dtypes(include=['string']).columns:
                df[col] = df[col].astype(object)
            try:
                source._connection.from_df(df).to_view(t)
            except duckdb.ParserException:
                continue
            new_tables[t] = source.sql_expr.format(table=t)
        source.tables = dict(new_tables, **tables)

        resolved_mirrors = {}
        for table, mirror in mirrors.items():
            if isinstance(mirror, tuple):
                src_spec, src_table = mirror
                source = cls.from_spec(src_spec)
                resolved_mirrors[table] = (source, src_table)
            elif mirror.get('type') == 'pipeline':
                from ..pipeline import Pipeline
                resolved_mirrors[table] = Pipeline.from_spec(mirror)
            else:
                resolved_mirrors[table] = Serializer.deserialize(mirror)
        spec['mirrors'] = resolved_mirrors
        return source

    def create_sql_expr_source(
        self, tables: dict[str, str], materialize: bool = True, params: dict[str, list | dict] | None = None, **kwargs
    ):
        """
        Creates a new SQL Source given a set of table names and
        corresponding SQL expressions.

        Arguments
        ---------
        tables: dict[str, str]
            Mapping from table name to SQL expression.
        materialize: bool
            Whether to materialize new tables
        params: dict[str, list | dict] | None
            Optional mapping from table name to parameters:
            - list: Positional parameters for ? placeholders
            - dict: Named parameters for $param_name placeholders
        kwargs: any
            Additional keyword arguments.

        Returns
        -------
        source: DuckDBSource
        """
        if params is None:
            params = {}

        source_params = dict(self.param.values(), **kwargs)

        # Only preserve existing tables if reusing the connection
        # If uri or initializers changed, start fresh with only new tables
        all_tables = tables
        if 'uri' not in kwargs and 'initializers' not in kwargs:
            # Reuse connection - start with ALL existing tables (upsert behavior)
            if isinstance(self.tables, dict):
                all_tables = dict(self.tables)
                # Update with new tables (overwrites if exists, adds if new)
                all_tables.update(tables)
            elif isinstance(self.tables, list):
                # Convert list to dict using sql_expr, then update with new tables
                all_tables = {t: self.sql_expr.format(table=t) for t in self.tables}
                all_tables.update(tables)
        else:
            # New connection - only use the new tables, but include file-based dependencies
            all_tables = dict(tables)
            # Analyze SQL expressions to find table dependencies
            for sql_expr in tables.values():
                if not isinstance(sql_expr, str):
                    continue
                try:
                    parsed = sqlglot.parse_one(sql_expr, dialect='duckdb')
                except Exception:
                    continue  # If parsing fails, continue without dependencies
                # Find all table references in the SQL
                # Add file-based tables that are referenced but not already included
                for table_obj in parsed.find_all(sqlglot.exp.Table):
                    table = table_obj.name
                    if table in self._file_based_tables and table not in all_tables:
                        all_tables[table] = self._file_based_tables[table]
        source_params['tables'] = all_tables

        if params:
            source_params['table_params'] = params
        # Reuse connection unless it has changed
        if 'uri' not in kwargs and 'initializers' not in kwargs:
            source_params['_connection'] = self._connection
        source_params.pop('name', None)
        source = type(self)(**source_params)

        if not materialize:
            return source

        for table, sql_expr in tables.copy().items():
            # Skip file paths - they're already handled by __init__
            if self._is_file_path(sql_expr):
                continue

            equivalent_sql_exprs = (
                self.sql_expr.format(table=f'"{table}"'),
                self.sql_expr.format(table=table),
            )
            if table in (self.tables or {}):
                # do not need to re-materialize existing
                continue
            elif sql_expr in equivalent_sql_exprs:
                continue
            table_expr = f'CREATE OR REPLACE TEMP TABLE "{table}" AS ({sql_expr})'
            cursor = self._connection.cursor()
            try:
                # Execute with parameters if provided for this table
                if table in params:
                    cursor.execute(table_expr, params[table])
                else:
                    cursor.execute(table_expr)
            except duckdb.CatalogException as e:
                original_e = e
                pattern = r"Table with name\s(\S+)"
                match = re.search(pattern, str(e))
                if match and isinstance(self.tables, dict):
                    name = match.group(1)
                    real = self.tables[name] if name in self.tables else self.tables[name.strip('"')]
                    table_expr = SQLSelectFrom(sql_expr=self.sql_expr, tables={name: real}).apply(table_expr)
                    try:
                        cursor.execute(table_expr)
                    except Exception as e:
                        raise original_e from e
                else:
                    raise e
            finally:
                cursor.close()

        # Preserve file-based metadata from parent source and merge with
        # any new file-based tables detected during __init__
        source._file_based_tables = {**self._file_based_tables, **source._file_based_tables}
        return source

    def execute(self, sql_query: str, params: list | dict | None = None, *args, **kwargs):
        with self._connection.cursor() as cursor:
            if params:
                return cursor.execute(sql_query, params, *args, **kwargs).fetch_df()
            return cursor.execute(sql_query, *args, **kwargs).fetch_df()

    def get_tables(self):
        if isinstance(self.tables, dict | list):
            return [t for t in list(self.tables) if not self._is_table_excluded(t)]

        with self._connection.cursor() as cursor:
            cursor.execute('SHOW TABLES')
            return [
                t[0] for t in cursor.fetchall()
                if not self._is_table_excluded(t[0])
            ]

    def normalize_table(self, table: str):
        tables = self.get_tables()
        if table not in tables and 'read_' in table:
            # Extract table name from read_* function
            table = re.search(r"read_(\w+)\('(.+?)'", table).group(2)
        table = normalize_table_name(table)
        return table

    @cached
    def get(self, table, **query):
        query.pop('__dask', None)
        sql_expr = self.get_sql_expr(table)
        sql_transforms = query.pop('sql_transforms', [])
        conditions = list(query.items())
        if self.filter_in_sql:
            sql_transforms = [SQLFilter(conditions=conditions)] + sql_transforms
        for st in sql_transforms:
            sql_expr = st.apply(sql_expr)

        # Apply stored SQL parameters if available for this table
        with self._connection.cursor() as cursor:
            if table in self.table_params:
                rel = cursor.execute(sql_expr, self.table_params[table])
            else:
                rel = cursor.execute(sql_expr)
            has_geom = any(d[0] == 'geometry' and d[1] == 'BINARY' for d in rel.description)
            df = rel.fetch_df(date_as_object=True)
            if has_geom:
                import geopandas as gpd
                geom_rel = cursor.execute(
                    f'SELECT ST_AsWKB(geometry::GEOMETRY) as geometry FROM ({sql_expr})'
                )
                geom_df = geom_rel.fetch_df()
                df['geometry'] = gpd.GeoSeries.from_wkb(geom_df.geometry.apply(bytes))
                df = gpd.GeoDataFrame(df)
        if not self.filter_in_sql:
            df = Filter.apply_to(df, conditions=conditions)
        return df

    def _get_table_metadata(self, tables: list[str]) -> dict[str, Any]:
        """
        Generate metadata for all tables or a single table (batched=False) in DuckDB.
        Handles formats: database.schema.table_name, schema.table_name, or table_name.

        Args:
            table: Table name(s) to get metadata for

        Returns:
            Dictionary with table metadata including description, columns, row count, etc.
        """
        metadata = {}
        for table_name in tables:
            sql_expr = self.get_sql_expr(table_name)
            schema_expr = SQLLimit(limit=0).apply(sql_expr)
            count_expr = SQLCount().apply(sql_expr)
            schema_result = self.execute(schema_expr)
            count = self.execute(count_expr).iloc[0, 0]
            table_metadata = {
                "description": "",
                "columns": {
                    col: {"data_type": str(dtype), "description": ""}
                    for col, dtype in zip(schema_result.columns, schema_result.dtypes, strict=False)
                },
                "rows": count,
                "updated_at": None,
                "created_at": None,
            }
            metadata[table_name] = table_metadata
        return metadata

    def close(self):
        """
        Close the DuckDB connection, releasing associated resources.

        This method should be called when the source is no longer needed to prevent
        connection leaks and properly clean up server-side resources.
        """
        if self._connection is not None:
            self._connection.close()
            self._connection = None
