from __future__ import annotations

import os
import re

from typing import TYPE_CHECKING, Any

import duckdb
import pandas as pd
import param

from ..config import config
from ..serializers import Serializer
from ..transforms import Filter
from ..transforms.sql import SQLFilter
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.""")

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

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

    source_type = 'duckdb'

    dialect = 'duckdb'

    def __init__(self, **params):
        connection = params.pop('_connection', None)
        super().__init__(**params)
        if connection:
            self._connection = connection
        else:
            self._connection = duckdb.connect(self.uri)
            for init in self.initializers:
                self._connection.execute(init)
        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):
                    self._connection.from_df(e.new).to_view(table)
                    self._set_cache(e.new, table)
                mirror.param.watch(update, 'data')
            try:
                self._connection.from_df(df).to_view(table)
            except (duckdb.CatalogException, duckdb.ParserException):
                continue

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

    @classmethod
    def _recursive_resolve(
        cls, spec: dict[str, Any], source_type: type[Source]
    ) -> tuple[dict[str, Any], dict[str, Any]]:
        mirrors = spec.pop('mirrors', {})
        spec, refs = super()._recursive_resolve(spec, source_type)
        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 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: str) -> str:
        # Look for read_* patterns like read_parquet, read_csv etc.
        matches = re.finditer(r"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():
                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():
            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 = spec['tables']
            spec['tables'] = {}
        else:
            ephemeral_tables = {}
        source = super().from_spec(spec)
        if not ephemeral_tables:
            return source
        new_tables = {}
        for t, data in ephemeral_tables.items():
            df = Serializer.deserialize(data)
            try:
                source._connection.from_df(df).to_view(t)
            except duckdb.ParserException:
                continue
            new_tables[t] = source.sql_expr.format(table=t)
        source.tables = new_tables
        return source

    def create_sql_expr_source(
        self, tables: dict[str, str], materialize: bool = True, **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
        kwargs: any
            Additional keyword arguments.

        Returns
        -------
        source: DuckDBSource
        """
        params = dict(self.param.values(), **kwargs)
        params['tables'] = tables
        # Reuse connection unless it has changed
        if 'uri' not in kwargs and 'initializers' not in kwargs:
            params['_connection'] = self._connection
        params.pop('name', None)
        source = type(self)(**params)
        if not materialize:
            return source

        for table, sql_expr in tables.copy().items():
            if sql_expr == self.sql_expr.format(table=f'"{table}"'):
                continue
            table_expr = f'CREATE OR REPLACE TEMP TABLE "{table}" AS ({sql_expr})'
            try:
                self._connection.execute(table_expr)
            except duckdb.CatalogException as 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('"')]
                    if 'select' not in real.lower() and not real.startswith('"'):
                        real = f'"{real}"'
                    else:
                        real = f'({real})'
                    table_expr = table_expr.replace(f'FROM {name}', f'FROM {real}')
                    source.tables[table] = sql_expr.replace(f'FROM {name}', f'FROM {real}')
                    self._connection.execute(table_expr)
                else:
                    raise e
        return source

    def execute(self, sql_query: str):
        return self._connection.execute(sql_query).fetch_df()

    def get_tables(self):
        if isinstance(self.tables, dict | list):
            return list(self.tables)
        return [t[0] for t in self._connection.execute('SHOW TABLES').fetchall()]

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

    def get_sql_expr(self, table: str):
        if isinstance(self.tables, dict):
            try:
                table = self.tables[self.normalize_table(table)]
            except KeyError:
                raise KeyError(f"Table {table} not found in {self.tables.keys()}")
        if '(' not in table and ')' not in table:
            table = f'"{table}"'

        # search if the so-called "table" already contains SELECT ... FROM
        # if so, we don't need to wrap it in a SELECT * FROM
        if re.search(r"(?i)\bselect\b[\s\S]+?\bfrom\b", table):
            sql_expr = table
        else:
            sql_expr = self.sql_expr.format(table=table)
        return sql_expr.rstrip(";")

    @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)
        rel = self._connection.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 = self._connection.execute(
                f'SELECT ST_AsWKB(geometry::GEOMETRY) as geometry FROM ({sql_expr})'
            ).fetch_df()
            df['geometry'] = gpd.GeoSeries.from_wkb(geom.geometry.apply(bytes))
            df = gpd.GeoDataFrame(df)
        if not self.filter_in_sql:
            df = Filter.apply_to(df, conditions=conditions)
        return df
