import asyncio
import json
import os
import urllib.request
import warnings
from typing import Literal
from datetime import datetime, timezone
from typing import Optional

warnings.filterwarnings(
    "ignore", message="Your application has authenticated using end user credentials"
)
import pandas as pd
import panel as pn
import param

def share(
    view,
    __name__,
    __file__,
    sharer: Literal["pn.serve", "embed.html", "pyodide"] | int | None = 2,
    pyodide_kwargs={"runtime": "pyodide-worker", "prerender": False, "inline": False},
    panel_serve_kwargs={"port": 8000, "show": True},
):

    if __name__ == "__main__":
        if isinstance(sharer, int):
            sharer = ["pn.serve", "embed.html", "pyodide", None][sharer]
        if pn.state._is_pyodide:
            return view.servable()
        elif sharer == "pyodide":
            import panel.io.convert
            import http.server
            import socketserver
            import webbrowser

            def start_http_server(port=8000, directory=".", show=True, fn=""):
                """
                Start an HTTP server to serve files from a specific directory.

                :param port: The port to serve on (default is 8000).
                :param directory: The directory to serve files from (default is the current directory).
                """
                # Change the working directory to the specified folder
                current_workdir = os.getcwd()
                os.chdir(directory)

                # Set up the HTTP server
                handler = http.server.SimpleHTTPRequestHandler
                with socketserver.TCPServer(("", port), handler) as httpd:
                    print(f"Serving HTTP on port {port} (http://localhost:{port}/) ...")
                    print(f"Serving files from directory: {os.path.abspath(directory)}")
                    try:
                        if show:
                            webbrowser.open(f"http://localhost:{port}/{fn}")
                        httpd.serve_forever()
                    except KeyboardInterrupt:
                        print("\nShutting down the server.")
                        httpd.server_close()
                        os.chdir(current_workdir)

            source_file = os.path.abspath(__file__)
            assert os.path.isfile(source_file), f"Can't find file `{__file__}`."
            workdir = os.path.join(os.path.split(__file__)[0], "dist")
            if not os.path.isdir(workdir):
                os.mkdir(workdir)
            _ = panel.io.convert.convert_app(
                source_file,
                dest_path=workdir,
                **pyodide_kwargs,
            )
            start_http_server(
                directory=workdir,
                fn=os.path.split(source_file)[1].replace(".py", ".html"),
            )
        elif sharer == "pn.serve":
            if "server" in globals():
                print("found it in globals")
                globals()["server"].stop()
            server: pn.io.server.Server = pn.serve(view, **panel_serve_kwargs)
            globals()["server"] = server
        elif sharer == "embed.html":
            fn = os.path.split(os.path.abspath(__file__))[-1].replace(".py", ".html")
            view.save(fn, embed=True, max_opts=50, embed_json=False)
        else:
            ...
    else:
        return view.servable()

if pn.state._is_pyodide:
    CACHE_TO_DISK = False
else:
    CACHE_TO_DISK = True
    import diskcache

BQ_BASE_URL = "https://bigquery.googleapis.com/bigquery/v2"
NUMERIC_TYPES = {"INTEGER", "INT64", "FLOAT", "FLOAT64", "NUMERIC", "BIGNUMERIC"}
TIMESTAMP_TYPES = {"TIMESTAMP", "DATETIME", "DATE"}


# @apply_to_all(log_class_methods, except_for=[])
class BigQuerier(param.Parameterized):

    job_project = param.String(precedence=-1)
    data_project = param.String(precedence=-1)
    dataset = param.String(precedence=-1)
    table = param.String(precedence=-1)
    access_token = param.String(default="", precedence=-1, doc="OAuth 2.0 access token. Required in Pyodide where google-auth is unavailable.")
    _access_token: str = ""  # class-level store for static method access

    columns = param.Dict(default={}, doc="Column name -> BigQuery type")
    ranges = param.Dict(default={}, doc="Numeric/timestamp column name -> (min, max)")
    data = param.DataFrame()
    is_loading = param.Boolean(
        default=False, doc="True while async requests are in flight"
    )
    nqueued = param.Integer(default=0, doc="Number of async requests currently queued")
    max_concurrent = param.Integer(
        default=5, bounds=(1, None), doc="Maximum number of concurrent async requests"
    )

    selected_columns = param.ListSelector(
        default=[], objects=set(), allow_None=True, doc="Selected column name"
    )
    indexer_columns = param.ListSelector(
        default=[], objects=set(), allow_None=True, doc="Selected column name"
    )

    def __init__(
        self,
        ignore_columns: list[str] | None = None,
        keep_columns: list[str] | None = None,
        categorical_columns: list[str] | None = None,
        **params,
    ):
        """Authenticate, fetch column schema and numeric ranges from the table.

        Parameters
        ----------
        ignore_columns : list[str], optional
            Column names to exclude from the schema entirely.
            Mutually exclusive with ``keep_columns``.
        keep_columns : list[str], optional
            Column names to keep; all other columns are ignored.
            Mutually exclusive with ``ignore_columns``.
        categorical_columns : list[str], optional
            Column names to treat as categorical (ListSelector) regardless of
            their BigQuery type.
        **params
            Keyword arguments passed to ``param.Parameterized``, including
            ``job_project``, ``data_project``, ``dataset``, and ``table``.
        """
        if ignore_columns is not None and keep_columns is not None:
            raise ValueError(
                "ignore_columns and keep_columns are mutually exclusive; "
                "specify one or the other, not both."
            )
        indexer_columns = params.pop("indexer_columns", [])
        selected_columns = params.pop("selected_columns", [])
        super().__init__(**params)
        if pn.state._is_pyodide and not self.access_token:
            raise ValueError("An access_token is required in Pyodide environment.")
        if self.access_token:
            BigQuerier._access_token = self.access_token

        self._indexer_params = set()
        self._indexer_watchers: dict[str, param.parameterized.Watcher] = {}
        self._categorical_columns: set[str] = set(categorical_columns or [])
        self._semaphore = asyncio.Semaphore(self.max_concurrent)
        self.columns = self._fetch_columns(
            ignore=ignore_columns or [], keep=keep_columns,
        )
        self.ranges = self._fetch_ranges()
        self.param["selected_columns"].objects = [None] + list(self.columns.keys())
        self.param["indexer_columns"].objects = [None] + list(self.columns.keys())

        self.param.update(
            selected_columns=selected_columns,
            indexer_columns=indexer_columns,
        )

        self.view = self.make_view()

    def make_view(self):

        indexer_widgets = list()
        for column in self.indexer_columns:
            column_type = self.columns.get(column)
            if column_type in NUMERIC_TYPES and column not in self._categorical_columns:
                widget = pn.widgets.RangeSlider.from_param(
                    self.param[f"indexer_{column}"],
                    name=column,
                )
            elif column_type in TIMESTAMP_TYPES and column not in self._categorical_columns:
                widget = pn.widgets.DateRangeSlider.from_param(
                    self.param[f"indexer_{column}"],
                    name=column,
                )
            else:
                widget = pn.widgets.MultiChoice.from_param(
                    self.param[f"indexer_{column}"],
                    name=column,
                )
            indexer_widgets.append(widget)

        return pn.Row(
            pn.Column(
                pn.widgets.MultiChoice.from_param(self.param["selected_columns"]),
                pn.widgets.MultiChoice.from_param(self.param["indexer_columns"]),
                pn.indicators.LoadingSpinner.from_param(
                    self.param["is_loading"],
                    name="Test",
                    size=22,
                    styles={"margin-left": "-31px"},
                ),
                pn.indicators.Number.from_param(self.param["nqueued"]),
                pn.widgets.Tabulator.from_param(self.param["data"], pagination="remote", page_size=10),
            ),
            pn.Column(
                *indexer_widgets
            ),
            sizing_mode="stretch_width",
        )

    @param.depends("indexer_columns", watch=True)
    def _on_indexer_columns_changed(self, param_name_format="indexer_{name}"):
        """Add/remove Range or DateRange parameters based on indexer column selection."""
        current = {c for c in self.indexer_columns if c is not None}
        previous = self._indexer_params

        # Remove parameters for deselected columns
        cls = type(self)
        for name in previous - current:
            formatted_name = param_name_format.format(name=name)
            if formatted_name in self._indexer_watchers:
                self.param.unwatch(self._indexer_watchers.pop(formatted_name))
            if hasattr(cls, formatted_name):
                delattr(cls, formatted_name)
            self._indexer_params.discard(name)

        # Add parameters for newly selected columns
        for name in current - previous:
            formatted_name = param_name_format.format(name=name)
            col_type = self.columns.get(name)
            bounds = self.ranges.get(name)
            if (
                col_type in NUMERIC_TYPES
                and bounds is not None
                and name not in self._categorical_columns
            ):
                self.param.add_parameter(
                    formatted_name, param.Range(default=bounds, bounds=bounds)
                )
            elif (
                col_type in TIMESTAMP_TYPES
                and bounds is not None
                and name not in self._categorical_columns
            ):
                self.param.add_parameter(
                    formatted_name, param.DateRange(default=bounds, bounds=bounds)
                )
            else:
                unique_vals = self._fetch_unique_values(name)
                self.param.add_parameter(
                    formatted_name,
                    param.ListSelector(default=[], objects=unique_vals),
                )
            self._indexer_watchers[formatted_name] = self.param.watch(
                self._on_indexer_value_changed, [formatted_name]
            )
            self._indexer_params.add(name)

    async def _on_indexer_value_changed(self, event: param.parameterized.Event) -> None:
        """Re-fetch selected columns when an indexer parameter value changes."""
        await self._on_selected_columns_changed()

    def _build_where_clause(
        self, param_name_format: str = "indexer_{name}"
    ) -> str | None:
        """Build a SQL WHERE clause from current indexer parameter values.

        Parameters
        ----------
        param_name_format : str
            Format string mapping column names to parameter names.

        Returns
        -------
        str or None
            A ``" WHERE ..."`` string, ``""`` if no filters are active,
            or ``None`` if the filters guarantee zero results.
        """
        conditions: list[str] = []
        for name in self._indexer_params:
            formatted_name = param_name_format.format(name=name)
            value = getattr(self, formatted_name, None)
            if value is None:
                continue
            col_type = self.columns.get(name)
            if col_type in NUMERIC_TYPES and name not in self._categorical_columns:
                lo, hi = value
                if lo is not None:
                    conditions.append(f"`{name}` >= {lo}")
                if hi is not None:
                    conditions.append(f"`{name}` <= {hi}")
            elif col_type in TIMESTAMP_TYPES and name not in self._categorical_columns:
                lo, hi = value
                if lo is not None:
                    conditions.append(f"`{name}` >= TIMESTAMP('{lo.isoformat()}')")
                if hi is not None:
                    conditions.append(f"`{name}` <= TIMESTAMP('{hi.isoformat()}')")
            else:
                if not value:
                    return None
                elif col_type in NUMERIC_TYPES:
                    joined = ", ".join(str(v) for v in value)
                    conditions.append(f"`{name}` IN ({joined})")
                else:
                    joined = ", ".join(f"'{v}'" for v in value)
                    conditions.append(f"`{name}` IN ({joined})")
        if not conditions:
            return ""
        return " WHERE " + " AND ".join(conditions)

    @param.depends("selected_columns", watch=True)
    async def _on_selected_columns_changed(self):
        """Fetch each selected column individually (for effective caching) and combine."""
        cols = [c for c in self.selected_columns if c is not None]
        if not cols:
            self.data = None
            return
        where = self._build_where_clause()
        if where is None:
            self.data = pd.DataFrame(columns=cols + list(self._indexer_params))
            return
        all_cols = cols + [c for c in self._indexer_params if c not in cols]
        dfs = await asyncio.gather(
            *(
                self.query(f"SELECT `{c}` FROM `{self.full_table_id}`{where}")
                for c in all_cols
            )
        )
        self.param.update(
            data = pd.concat(dfs, axis=1)
        )

    @property
    def full_table_id(self) -> str:
        """Fully qualified BigQuery table ID: ``data_project.dataset.table``."""
        return f"{self.data_project}.{self.dataset}.{self.table}"

    @staticmethod
    @pn.cache(
        to_disk=CACHE_TO_DISK,
        cache_path="./big_querier_cache",
        per_session=False,
        ttl=5 * 86400,
    )
    def _request(method: str, url: str, body: Optional[dict] = None) -> dict:
        """Make a synchronous HTTP request to the BigQuery REST API.

        Token is fetched internally so it does not affect the cache key.

        Parameters
        ----------
        method : str
            HTTP method (e.g. ``"GET"``, ``"POST"``).
        url : str
            Full request URL.
        body : dict, optional
            JSON-serializable request body.

        Returns
        -------
        dict
            Parsed JSON response.
        """
        print(f"Running {body}")
        token = BigQuerier._get_access_token()
        headers = {
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json",
        }
        data = json.dumps(body).encode() if body else None
        req = urllib.request.Request(url, data=data, headers=headers, method=method)
        with urllib.request.urlopen(req) as resp:
            return json.loads(resp.read())

    @staticmethod
    @pn.cache(
        to_disk=CACHE_TO_DISK, cache_path="./big_querier_cache", per_session=False
    )
    async def _async_request(
        method: str, url: str, body: Optional[dict] = None
    ) -> dict:
        """Make an asynchronous HTTP request to the BigQuery REST API.

        Token is fetched internally so it does not affect the cache key.

        Parameters
        ----------
        method : str
            HTTP method (e.g. ``"GET"``, ``"POST"``).
        url : str
            Full request URL.
        body : dict, optional
            JSON-serializable request body.

        Returns
        -------
        dict
            Parsed JSON response.
        """
        print(f"Running {body}")
        if pn.state._is_pyodide:
            token = BigQuerier._get_access_token()
        else:
            token = await asyncio.to_thread(BigQuerier._get_access_token)
        headers = {"Authorization": f"Bearer {token}"}
        if pn.state._is_pyodide:
            from pyodide.http import (  # pyright: ignore [reportMissingImports] # noqa
                pyfetch,
            )

            kwargs = {"method": method, "headers": headers}
            if body is not None:
                kwargs["body"] = json.dumps(body)
            resp = await pyfetch(url, **kwargs)
            return await resp.json()
        else:
            import aiohttp
            async with aiohttp.ClientSession() as session:
                req = getattr(session, method.lower())
                async with req(url, json=body, headers=headers) as resp:
                    resp.raise_for_status()
                    return await resp.json()

    def _fetch_columns(
        self,
        ignore: list[str] | None = None,
        keep: list[str] | None = None,
    ) -> dict[str, str]:
        """Fetch column names and BigQuery types via the tables.get API.

        Parameters
        ----------
        ignore : list[str], optional
            Column names to exclude from the result.
        keep : list[str], optional
            Column names to include; all others are excluded.

        Returns
        -------
        dict[str, str]
            Mapping of column name to BigQuery type string (e.g. ``"FLOAT"``).
        """
        skip = set(ignore or [])
        allow = set(keep) if keep is not None else None
        url = (
            f"{BQ_BASE_URL}/projects/{self.data_project}"
            f"/datasets/{self.dataset}/tables/{self.table}"
        )
        response = self._request("GET", url)
        return {
            f["name"]: f["type"]
            for f in response["schema"]["fields"]
            if f["name"] not in skip
            and (allow is None or f["name"] in allow)
        }

    def _fetch_ranges(
        self,
    ) -> dict[
        str, tuple[float | None, float | None] | tuple[datetime | None, datetime | None]
    ]:
        """Fetch (min, max) ranges for numeric and timestamp columns via a SQL query.

        Returns
        -------
        dict[str, tuple]
            Mapping of column name to ``(min, max)`` tuple. Values are
            ``float`` for numeric columns and ``datetime`` for timestamp columns.
        """
        numeric_cols = [
            name for name, dtype in self.columns.items() if dtype in NUMERIC_TYPES
        ]
        timestamp_cols = [
            name for name, dtype in self.columns.items() if dtype in TIMESTAMP_TYPES
        ]
        range_cols = numeric_cols + timestamp_cols
        if not range_cols:
            return {}

        agg_parts = []
        for col in range_cols:
            agg_parts.append(f"MIN(`{col}`) AS `{col}__min`")
            agg_parts.append(f"MAX(`{col}`) AS `{col}__max`")

        query = f"SELECT {', '.join(agg_parts)} FROM `{self.full_table_id}`"
        url = f"{BQ_BASE_URL}/projects/{self.job_project}/queries"
        response = self._request("POST", url, {"query": query, "useLegacySql": False})

        row = response["rows"][0]["f"]
        ranges = {}
        for i, col in enumerate(range_cols):
            min_val = row[i * 2]["v"]
            max_val = row[i * 2 + 1]["v"]
            if col in timestamp_cols:
                ranges[col] = (
                    (
                        datetime.fromtimestamp(float(min_val), tz=timezone.utc)
                        if min_val is not None
                        else None
                    ),
                    (
                        datetime.fromtimestamp(float(max_val), tz=timezone.utc)
                        if max_val is not None
                        else None
                    ),
                )
            else:
                ranges[col] = (
                    float(min_val) if min_val is not None else None,
                    float(max_val) if max_val is not None else None,
                )
        return ranges

    def _fetch_unique_values(self, column: str) -> list[str]:
        """Fetch unique values for a column via a synchronous SQL query.

        Parameters
        ----------
        column : str
            Column name to query distinct values for.

        Returns
        -------
        list[str]
            Sorted list of unique values.
        """
        self.nqueued += 1
        self.is_loading = True
        try:
            query = f"SELECT DISTINCT `{column}` FROM `{self.full_table_id}` ORDER BY `{column}`"
            url = f"{BQ_BASE_URL}/projects/{self.job_project}/queries"
            response = BigQuerier._request(
                "POST", url, {"query": query, "useLegacySql": False}
            )
            if "rows" not in response:
                return []
            return [
                row["f"][0]["v"]
                for row in response["rows"]
                if row["f"][0]["v"] is not None
            ]
        finally:
            self.nqueued -= 1
            self.is_loading = self.nqueued > 0

    async def _async_fetch_unique_values(self, column: str) -> list[str]:
        """Fetch unique values for a column via an async SQL query.

        Parameters
        ----------
        column : str
            Column name to query distinct values for.

        Returns
        -------
        list[str]
            Sorted list of unique values.
        """
        self.nqueued += 1
        self.is_loading = True
        try:
            async with self._semaphore:
                query = f"SELECT DISTINCT `{column}` FROM `{self.full_table_id}` ORDER BY `{column}`"
                url = f"{BQ_BASE_URL}/projects/{self.job_project}/queries"
                response = await self._async_request(
                    "POST", url, {"query": query, "useLegacySql": False}
                )
            if "rows" not in response:
                return []
            return [
                row["f"][0]["v"]
                for row in response["rows"]
                if row["f"][0]["v"] is not None
            ]
        finally:
            self.nqueued -= 1
            self.is_loading = self.nqueued > 0

    @staticmethod
    def _get_access_token() -> str:
        """Get an access token using Application Default Credentials.

        In Pyodide, supply a token via ``BigQuerier(access_token="...")``.

        Returns
        -------
        str
            OAuth 2.0 bearer token.
        """
        if BigQuerier._access_token:
            return BigQuerier._access_token
        import google.auth
        import google.auth.transport.requests
        credentials, _ = google.auth.default(
            scopes=["https://www.googleapis.com/auth/bigquery"]
        )
        credentials.refresh(google.auth.transport.requests.Request())
        return credentials.token

    @staticmethod
    def _rows_to_dataframe(response: dict) -> pd.DataFrame:
        """Convert a BigQuery query response JSON to a pandas DataFrame.

        Parameters
        ----------
        response : dict
            Raw JSON response from the BigQuery queries API.

        Returns
        -------
        pd.DataFrame
            DataFrame with column names from the response schema.
        """
        fields = response["schema"]["fields"]
        col_names = [f["name"] for f in fields]

        if "rows" not in response:
            return pd.DataFrame(columns=col_names)

        rows = [[cell["v"] for cell in row["f"]] for row in response["rows"]]
        df = pd.DataFrame(rows, columns=col_names)
        for field in fields:
            col = field["name"]
            if field["type"] in TIMESTAMP_TYPES:
                df[col] = pd.to_datetime(df[col].astype(float), unit="s")
            elif field["type"] in NUMERIC_TYPES:
                df[col] = pd.to_numeric(df[col])
        return df

    async def query(self, query: str) -> pd.DataFrame:
        """Execute an async SQL query and return the results as a DataFrame.

        Parameters
        ----------
        query : str
            SQL query string (Standard SQL).

        Returns
        -------
        pd.DataFrame
            Query results. Also stored in ``self.data``.
        """
        self.nqueued += 1
        self.is_loading = True
        try:
            async with self._semaphore:
                url = f"{BQ_BASE_URL}/projects/{self.job_project}/queries"
                data = await self._async_request(
                    "POST", url, {"query": query, "useLegacySql": False}
                )
            data_ = self._rows_to_dataframe(data)
            return data_
        finally:
            self.nqueued -= 1
            self.is_loading = self.nqueued > 0
