import duckdb
import pandas as pd
from typing import Union, List, Optional
from .protocols import (
    ChartBuilderProtocol, 
    QueryBuilderProtocol,
    DatasetProtocol
)

from .chart_builders import get_chart_builder
from .query_builders import get_query_builder

class DuckViz:
    def __init__(
        self,
        dataset : DatasetProtocol,
        chart_builder: str = 'plotly',
        query_builder: str = 'duckdb',
        theme: Optional[str] = None,
    ):
        self.dataset = dataset
        self.df = self.dataset.get_data()
        self.chart_builder : ChartBuilderProtocol = get_chart_builder(chart_builder)
        self.query_builder : QueryBuilderProtocol = get_query_builder(query_builder)
        self.theme = theme

        self.conn = duckdb.connect(database=":memory:")
        self.conn.register("data", self.df)

        self._where_clause = None

    def _execute(self, query: str) -> pd.DataFrame:
        return self.conn.execute(query).df()

    def filter(self, condition: str):
        """
        Add SQL WHERE clause.
        Example:
            dv.filter("region = 'West'")
        """
        self._where_clause = condition
        return self

    def bar_chart(
        self,
        x_axis: Optional[str],
        y_axis: Union[str, List[str]],
        aggregation: Optional[Union[str, List[str]]] = "sum",
        title: Optional[str] = None,
        x_axis_title: Optional[str] = None,
        y_axis_title: Optional[str] = None,
    ):

        result = self.query_builder.build(
            x=x_axis,
            y=y_axis,
            agg=aggregation,
            where=self._where_clause,
            group=True,
        )

        df_result = self._execute(result["query"])

        return self.chart_builder.bar(
            df=df_result,
            x_axis=result["x_alias"],
            y_axes=result["y_aliases"],
            title=title,
            x_axis_title=x_axis_title,
            y_axis_title=y_axis_title,
            theme=self.theme,
        )

    def line_chart(
        self,
        x_axis: Optional[str],
        y_axis: Union[str, List[str]],
        aggregation: Optional[Union[str, List[str]]] = "sum",
        title: Optional[str] = None,
        x_axis_title: Optional[str] = None,
        y_axis_title: Optional[str] = None,
    ):

        result = self.query_builder.build(
            x=x_axis,
            y=y_axis,
            agg=aggregation,
            where=self._where_clause,
            group=True,
        )

        df_result = self._execute(result["query"])

        return self.chart_builder.line(
            df=df_result,
            x_axis=result["x_alias"],
            y_axes=result["y_aliases"],
            title=title,
            x_axis_title=x_axis_title,
            y_axis_title=y_axis_title,
            theme=self.theme,
        )

    def area_chart(
        self,
        x_axis: Optional[str],
        y_axis: Union[str, List[str]],
        aggregation: Optional[Union[str, List[str]]] = "sum",
        title: Optional[str] = None,
        x_axis_title: Optional[str] = None,
        y_axis_title: Optional[str] = None,
    ):

        result = self.query_builder.build(
            x=x_axis,
            y=y_axis,
            agg=aggregation,
            where=self._where_clause,
            group=True,
        )

        df_result = self._execute(result["query"])

        return self.chart_builder.area(
            df=df_result,
            x_axis=result["x_alias"],
            y_axes=result["y_aliases"],
            title=title,
            x_axis_title=x_axis_title,
            y_axis_title=y_axis_title,
            theme=self.theme,
        )

    def pie_chart(
        self,
        names: str,
        values: str,
        aggregation: Optional[str] = "sum",
        title: Optional[str] = None,
    ):

        result = self.query_builder.build(
            x=names,
            y=values,
            agg=aggregation,
            where=self._where_clause,
            group=True,
        )

        df_result = self._execute(result["query"])

        return self.chart_builder.pie(
            df=df_result,
            names=result["x_alias"],
            values=result["y_aliases"][0],
            title=title,
            theme=self.theme,
        )

    def histogram(
        self,
        column: str,
        bins: int = 20,
        title: Optional[str] = None,
        x_axis_title: Optional[str] = None,
        y_axis_title: Optional[str] = None,
    ):

        result = self.query_builder.build(
            x=None,
            y=column,
            agg=None,
            where=self._where_clause,
            group=False,
        )

        df_result = self._execute(result["query"])

        return self.chart_builder.histogram(
            df=df_result,
            column=result["y_aliases"][0],
            bins=bins,
            title=title,
            x_axis_title=x_axis_title,
            y_axis_title=y_axis_title,
            theme=self.theme,
        )

    def combo_chart(
        self,
        x_axis: str,
        y1: str,
        y2: str,
        agg1: Optional[str] = "sum",
        agg2: Optional[str] = "sum",
        title: Optional[str] = None,
    ):

        result = self.query_builder.build(
            x=x_axis,
            y=[y1, y2],
            agg=[agg1, agg2],
            where=self._where_clause,
            group=True,
        )

        df_result = self._execute(result["query"])

        return self.chart_builder.combo(
            df=df_result,
            x_axis=result["x_alias"],
            y1=result["y_aliases"][0],
            y2=result["y_aliases"][1],
            label1=y1,
            label2=y2,
            title=title,
            theme=self.theme,
        )
    
    def kpi(self, metric, aggregation="sum", title=None):

        result = self.query_builder.build_kpi(
            metric=metric,
            aggregation=aggregation,
            where=self._where_clause
        )

        df_result = self._execute(result["query"])
        value = df_result.iloc[0, 0]

        return self.chart_builder.kpi(
            value=value,
            title=title,
            theme=self.theme,
        )
    
    def kpi_compare(
        self,
        metric,
        date_column,
        aggregation="sum",
        period="month",
        title=None,
    ):

        result = self.query_builder.build_kpi_compare(
            metric=metric,
            date_column=date_column,
            aggregation=aggregation,
            period=period,
            where=self._where_clause
        )

        df_result = self._execute(result["query"])

        if len(df_result) < 2:
            current = df_result.iloc[0]["value"]
            previous = None
            change_pct = None
        else:
            current = df_result.iloc[0]["value"]
            previous = df_result.iloc[1]["value"]
            change_pct = ((current - previous) / previous) * 100 if previous else None

        return self.chart_builder.kpi_compare(
            current=current,
            previous=previous,
            change_pct=change_pct,
            title=title,
            theme=self.theme,
        )


