# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import functools
import json
import types
from typing import Callable
from typing import Optional

from google.auth.credentials import Credentials
from google.cloud import bigquery

from . import client
from ..tool_context import ToolContext
from .config import BigQueryToolConfig
from .config import WriteMode

BIGQUERY_SESSION_INFO_KEY = "bigquery_session_info"


def execute_sql(
    project_id: str,
    query: str,
    credentials: Credentials,
    settings: BigQueryToolConfig,
    tool_context: ToolContext,
) -> dict:
  """Run a BigQuery or BigQuery ML SQL query in the project and return the result.

  Args:
      project_id (str): The GCP project id in which the query should be
        executed.
      query (str): The BigQuery SQL query to be executed.
      credentials (Credentials): The credentials to use for the request.
      settings (BigQueryToolConfig): The settings for the tool.
      tool_context (ToolContext): The context for the tool.

  Returns:
      dict: Dictionary representing the result of the query.
            If the result contains the key "result_is_likely_truncated" with
            value True, it means that there may be additional rows matching the
            query not returned in the result.

  Examples:
      Fetch data or insights from a table:

          >>> execute_sql("my_project",
          ... "SELECT island, COUNT(*) AS population "
          ... "FROM bigquery-public-data.ml_datasets.penguins GROUP BY island")
          {
            "status": "SUCCESS",
            "rows": [
                {
                    "island": "Dream",
                    "population": 124
                },
                {
                    "island": "Biscoe",
                    "population": 168
                },
                {
                    "island": "Torgersen",
                    "population": 52
                }
            ]
          }
  """
  try:
    # Validate compute project if applicable
    if (
        settings.compute_project_id
        and project_id != settings.compute_project_id
    ):
      return {
          "status": "ERROR",
          "error_details": (
              f"Cannot execute query in the project {project_id}, as the tool"
              " is restricted to execute queries only in the project"
              f" {settings.compute_project_id}."
          ),
      }

    # Get BigQuery client
    bq_client = client.get_bigquery_client(
        project=project_id,
        credentials=credentials,
        location=settings.location,
        user_agent=settings.application_name,
    )

    # BigQuery connection properties where applicable
    bq_connection_properties = None

    if not settings or settings.write_mode == WriteMode.BLOCKED:
      dry_run_query_job = bq_client.query(
          query,
          project=project_id,
          job_config=bigquery.QueryJobConfig(dry_run=True),
      )
      if dry_run_query_job.statement_type != "SELECT":
        return {
            "status": "ERROR",
            "error_details": "Read-only mode only supports SELECT statements.",
        }
    elif settings.write_mode == WriteMode.PROTECTED:
      # In protected write mode, write operation only to a temporary artifact is
      # allowed. This artifact must have been created in a BigQuery session. In
      # such a scenario the session info (session id and the anonymous dataset
      # containing the artifact) is persisted in the tool context.
      bq_session_info = tool_context.state.get(BIGQUERY_SESSION_INFO_KEY, None)
      if bq_session_info:
        bq_session_id, bq_session_dataset_id = bq_session_info
      else:
        session_creator_job = bq_client.query(
            "SELECT 1",
            project=project_id,
            job_config=bigquery.QueryJobConfig(
                dry_run=True, create_session=True
            ),
        )
        bq_session_id = session_creator_job.session_info.session_id
        bq_session_dataset_id = session_creator_job.destination.dataset_id

        # Remember the BigQuery session info for subsequent queries
        tool_context.state[BIGQUERY_SESSION_INFO_KEY] = (
            bq_session_id,
            bq_session_dataset_id,
        )

      # Session connection property will be set in the query execution
      bq_connection_properties = [
          bigquery.ConnectionProperty("session_id", bq_session_id)
      ]

      # Check the query type w.r.t. the BigQuery session
      dry_run_query_job = bq_client.query(
          query,
          project=project_id,
          job_config=bigquery.QueryJobConfig(
              dry_run=True,
              connection_properties=bq_connection_properties,
          ),
      )
      if (
          dry_run_query_job.statement_type != "SELECT"
          and dry_run_query_job.destination.dataset_id != bq_session_dataset_id
      ):
        return {
            "status": "ERROR",
            "error_details": (
                "Protected write mode only supports SELECT statements, or write"
                " operations in the anonymous dataset of a BigQuery session."
            ),
        }

    # Finally execute the query and fetch the result
    job_config = (
        bigquery.QueryJobConfig(connection_properties=bq_connection_properties)
        if bq_connection_properties
        else None
    )
    row_iterator = bq_client.query_and_wait(
        query,
        job_config=job_config,
        project=project_id,
        max_results=settings.max_query_result_rows,
    )
    rows = []
    for row in row_iterator:
      row_values = {}
      for key, val in row.items():
        try:
          # if the json serialization of the value succeeds, use it as is
          json.dumps(val)
        except:
          val = str(val)
        row_values[key] = val
      rows.append(row_values)

    result = {"status": "SUCCESS", "rows": rows}
    if (
        settings.max_query_result_rows is not None
        and len(rows) == settings.max_query_result_rows
    ):
      result["result_is_likely_truncated"] = True
    return result
  except Exception as ex:  # pylint: disable=broad-except
    return {
        "status": "ERROR",
        "error_details": str(ex),
    }


def _execute_sql_write_mode(*args, **kwargs) -> dict:
  """Run a BigQuery or BigQuery ML SQL query in the project and return the result.

  Args:
      project_id (str): The GCP project id in which the query should be
        executed.
      query (str): The BigQuery SQL query to be executed.
      credentials (Credentials): The credentials to use for the request.
      settings (BigQueryToolConfig): The settings for the tool.
      tool_context (ToolContext): The context for the tool.

  Returns:
      dict: Dictionary representing the result of the query.
            If the result contains the key "result_is_likely_truncated" with
            value True, it means that there may be additional rows matching the
            query not returned in the result.

  Examples:
      Fetch data or insights from a table:

          >>> execute_sql("my_project",
          ... "SELECT island, COUNT(*) AS population "
          ... "FROM bigquery-public-data.ml_datasets.penguins GROUP BY island")
          {
            "status": "SUCCESS",
            "rows": [
                {
                    "island": "Dream",
                    "population": 124
                },
                {
                    "island": "Biscoe",
                    "population": 168
                },
                {
                    "island": "Torgersen",
                    "population": 52
                }
            ]
          }

      Create a table with schema prescribed:

          >>> execute_sql("my_project",
          ... "CREATE TABLE my_project.my_dataset.my_table "
          ... "(island STRING, population INT64)")
          {
            "status": "SUCCESS",
            "rows": []
          }

      Insert data into an existing table:

          >>> execute_sql("my_project",
          ... "INSERT INTO my_project.my_dataset.my_table (island, population) "
          ... "VALUES ('Dream', 124), ('Biscoe', 168)")
          {
            "status": "SUCCESS",
            "rows": []
          }

      Create a table from the result of a query:

          >>> execute_sql("my_project",
          ... "CREATE TABLE my_project.my_dataset.my_table AS "
          ... "SELECT island, COUNT(*) AS population "
          ... "FROM bigquery-public-data.ml_datasets.penguins GROUP BY island")
          {
            "status": "SUCCESS",
            "rows": []
          }

      Delete a table:

          >>> execute_sql("my_project",
          ... "DROP TABLE my_project.my_dataset.my_table")
          {
            "status": "SUCCESS",
            "rows": []
          }

      Copy a table to another table:

          >>> execute_sql("my_project",
          ... "CREATE TABLE my_project.my_dataset.my_table_clone "
          ... "CLONE my_project.my_dataset.my_table")
          {
            "status": "SUCCESS",
            "rows": []
          }

      Create a snapshot (a lightweight, read-optimized copy) of en existing
      table:

          >>> execute_sql("my_project",
          ... "CREATE SNAPSHOT TABLE my_project.my_dataset.my_table_snapshot "
          ... "CLONE my_project.my_dataset.my_table")
          {
            "status": "SUCCESS",
            "rows": []
          }

      Create a BigQuery ML linear regression model:

          >>> execute_sql("my_project",
          ... "CREATE MODEL `my_dataset.my_model` "
          ... "OPTIONS (model_type='linear_reg', input_label_cols=['body_mass_g']) AS "
          ... "SELECT * FROM `bigquery-public-data.ml_datasets.penguins` "
          ... "WHERE body_mass_g IS NOT NULL")
          {
            "status": "SUCCESS",
            "rows": []
          }

      Evaluate BigQuery ML model:

          >>> execute_sql("my_project",
          ... "SELECT * FROM ML.EVALUATE(MODEL `my_dataset.my_model`)")
          {
            "status": "SUCCESS",
            "rows": [{'mean_absolute_error': 227.01223667447218,
                      'mean_squared_error': 81838.15989216768,
                      'mean_squared_log_error': 0.0050704473735013,
                      'median_absolute_error': 173.08081641661738,
                      'r2_score': 0.8723772534253441,
                      'explained_variance': 0.8723772534253442}]
          }

      Evaluate BigQuery ML model on custom data:

          >>> execute_sql("my_project",
          ... "SELECT * FROM ML.EVALUATE(MODEL `my_dataset.my_model`, "
          ... "(SELECT * FROM `my_dataset.my_table`))")
          {
            "status": "SUCCESS",
            "rows": [{'mean_absolute_error': 227.01223667447218,
                      'mean_squared_error': 81838.15989216768,
                      'mean_squared_log_error': 0.0050704473735013,
                      'median_absolute_error': 173.08081641661738,
                      'r2_score': 0.8723772534253441,
                      'explained_variance': 0.8723772534253442}]
          }

      Predict using BigQuery ML model:

          >>> execute_sql("my_project",
          ... "SELECT * FROM ML.PREDICT(MODEL `my_dataset.my_model`, "
          ... "(SELECT * FROM `my_dataset.my_table`))")
          {
            "status": "SUCCESS",
            "rows": [
                {
                  "predicted_body_mass_g": "3380.9271650847013",
                  ...
                }, {
                  "predicted_body_mass_g": "3873.6072435386004",
                  ...
                },
                ...
            ]
          }

      Delete a BigQuery ML model:

          >>> execute_sql("my_project", "DROP MODEL `my_dataset.my_model`")
          {
            "status": "SUCCESS",
            "rows": []
          }

  Notes:
      - If a destination table already exists, there are a few ways to overwrite
      it:
          - Use "CREATE OR REPLACE TABLE" instead of "CREATE TABLE".
          - First run "DROP TABLE", followed by "CREATE TABLE".
      - If a model already exists, there are a few ways to overwrite it:
          - Use "CREATE OR REPLACE MODEL" instead of "CREATE MODEL".
          - First run "DROP MODEL", followed by "CREATE MODEL".
  """
  return execute_sql(*args, **kwargs)


def _execute_sql_protected_write_mode(*args, **kwargs) -> dict:
  """Run a BigQuery or BigQuery ML SQL query in the project and return the result.

  Args:
      project_id (str): The GCP project id in which the query should be
        executed.
      query (str): The BigQuery SQL query to be executed.
      credentials (Credentials): The credentials to use for the request.
      settings (BigQueryToolConfig): The settings for the tool.
      tool_context (ToolContext): The context for the tool.

  Returns:
      dict: Dictionary representing the result of the query.
            If the result contains the key "result_is_likely_truncated" with
            value True, it means that there may be additional rows matching the
            query not returned in the result.

  Examples:
      Fetch data or insights from a table:

          >>> execute_sql("my_project",
          ... "SELECT island, COUNT(*) AS population "
          ... "FROM bigquery-public-data.ml_datasets.penguins GROUP BY island")
          {
            "status": "SUCCESS",
            "rows": [
                {
                    "island": "Dream",
                    "population": 124
                },
                {
                    "island": "Biscoe",
                    "population": 168
                },
                {
                    "island": "Torgersen",
                    "population": 52
                }
            ]
          }

      Create a temporary table with schema prescribed:

          >>> execute_sql("my_project",
          ... "CREATE TEMP TABLE my_table (island STRING, population INT64)")
          {
            "status": "SUCCESS",
            "rows": []
          }

      Insert data into an existing temporary table:

          >>> execute_sql("my_project",
          ... "INSERT INTO my_table (island, population) "
          ... "VALUES ('Dream', 124), ('Biscoe', 168)")
          {
            "status": "SUCCESS",
            "rows": []
          }

      Create a temporary table from the result of a query:

          >>> execute_sql("my_project",
          ... "CREATE TEMP TABLE my_table AS "
          ... "SELECT island, COUNT(*) AS population "
          ... "FROM bigquery-public-data.ml_datasets.penguins GROUP BY island")
          {
            "status": "SUCCESS",
            "rows": []
          }

      Delete a temporary table:

          >>> execute_sql("my_project", "DROP TABLE my_table")
          {
            "status": "SUCCESS",
            "rows": []
          }

      Copy a temporary table to another temporary table:

          >>> execute_sql("my_project",
          ... "CREATE TEMP TABLE my_table_clone CLONE my_table")
          {
            "status": "SUCCESS",
            "rows": []
          }

      Create a temporary BigQuery ML linear regression model:

          >>> execute_sql("my_project",
          ... "CREATE TEMP MODEL my_model "
          ... "OPTIONS (model_type='linear_reg', input_label_cols=['body_mass_g']) AS"
          ... "SELECT * FROM `bigquery-public-data.ml_datasets.penguins` "
          ... "WHERE body_mass_g IS NOT NULL")
          {
            "status": "SUCCESS",
            "rows": []
          }

      Evaluate BigQuery ML model:

          >>> execute_sql("my_project", "SELECT * FROM ML.EVALUATE(MODEL my_model)")
          {
            "status": "SUCCESS",
            "rows": [{'mean_absolute_error': 227.01223667447218,
                      'mean_squared_error': 81838.15989216768,
                      'mean_squared_log_error': 0.0050704473735013,
                      'median_absolute_error': 173.08081641661738,
                      'r2_score': 0.8723772534253441,
                      'explained_variance': 0.8723772534253442}]
          }

      Evaluate BigQuery ML model on custom data:

          >>> execute_sql("my_project",
          ... "SELECT * FROM ML.EVALUATE(MODEL my_model, "
          ... "(SELECT * FROM `my_dataset.my_table`))")
          {
            "status": "SUCCESS",
            "rows": [{'mean_absolute_error': 227.01223667447218,
                      'mean_squared_error': 81838.15989216768,
                      'mean_squared_log_error': 0.0050704473735013,
                      'median_absolute_error': 173.08081641661738,
                      'r2_score': 0.8723772534253441,
                      'explained_variance': 0.8723772534253442}]
          }

      Predict using BigQuery ML model:

          >>> execute_sql("my_project",
          ... "SELECT * FROM ML.PREDICT(MODEL my_model, "
          ... "(SELECT * FROM `my_dataset.my_table`))")
          {
            "status": "SUCCESS",
            "rows": [
                {
                  "predicted_body_mass_g": "3380.9271650847013",
                  ...
                }, {
                  "predicted_body_mass_g": "3873.6072435386004",
                  ...
                },
                ...
            ]
          }

      Delete a BigQuery ML model:

          >>> execute_sql("my_project", "DROP MODEL my_model")
          {
            "status": "SUCCESS",
            "rows": []
          }

  Notes:
      - If a destination table already exists, there are a few ways to overwrite
      it:
          - Use "CREATE OR REPLACE TEMP TABLE" instead of "CREATE TEMP TABLE".
          - First run "DROP TABLE", followed by "CREATE TEMP TABLE".
      - Only temporary tables can be created, inserted into or deleted. Please
      do not try creating a permanent table (non-TEMP table), inserting into or
      deleting one.
      - If a destination model already exists, there are a few ways to overwrite
      it:
          - Use "CREATE OR REPLACE TEMP MODEL" instead of "CREATE TEMP MODEL".
          - First run "DROP MODEL", followed by "CREATE TEMP MODEL".
      - Only temporary models can be created or deleted. Please do not try
      creating a permanent model (non-TEMP model) or deleting one.
  """
  return execute_sql(*args, **kwargs)


def get_execute_sql(settings: BigQueryToolConfig) -> Callable[..., dict]:
  """Get the execute_sql tool customized as per the given tool settings.

  Args:
      settings: BigQuery tool settings indicating the behavior of the
        execute_sql tool.

  Returns:
      callable[..., dict]: A version of the execute_sql tool respecting the tool
      settings.
  """

  if not settings or settings.write_mode == WriteMode.BLOCKED:
    return execute_sql

  # Create a new function object using the original function's code and globals.
  # We pass the original code, globals, name, defaults, and closure.
  # This creates a raw function object without copying other metadata yet.
  execute_sql_wrapper = types.FunctionType(
      execute_sql.__code__,
      execute_sql.__globals__,
      execute_sql.__name__,
      execute_sql.__defaults__,
      execute_sql.__closure__,
  )

  # Use functools.update_wrapper to copy over other essential attributes
  # from the original function to the new one.
  # This includes __name__, __qualname__, __module__, __annotations__, etc.
  # It specifically allows us to then set __doc__ separately.
  functools.update_wrapper(execute_sql_wrapper, execute_sql)

  # Now, set the new docstring
  if settings.write_mode == WriteMode.PROTECTED:
    execute_sql_wrapper.__doc__ = _execute_sql_protected_write_mode.__doc__
  else:
    execute_sql_wrapper.__doc__ = _execute_sql_write_mode.__doc__

  return execute_sql_wrapper


def forecast(
    project_id: str,
    history_data: str,
    timestamp_col: str,
    data_col: str,
    horizon: int = 10,
    id_cols: Optional[list[str]] = None,
    *,
    credentials: Credentials,
    settings: BigQueryToolConfig,
    tool_context: ToolContext,
) -> dict:
  """Run a BigQuery AI time series forecast using AI.FORECAST.

  Args:
      project_id (str): The GCP project id in which the query should be
        executed.
      history_data (str): The table id of the BigQuery table containing the
        history time series data or a query statement that select the history
        data.
      timestamp_col (str): The name of the column containing the timestamp for
        each data point.
      data_col (str): The name of the column containing the numerical values to
        be forecasted.
      horizon (int, optional): The number of time steps to forecast into the
        future. Defaults to 10.
      id_cols (list, optional): The column names of the id columns to indicate
        each time series when there are multiple time series in the table. All
        elements must be strings. Defaults to None.
      credentials (Credentials): The credentials to use for the request.
      settings (BigQueryToolConfig): The settings for the tool.
      tool_context (ToolContext): The context for the tool.

  Returns:
      dict: Dictionary representing the result of the forecast. The result
            contains the forecasted values along with prediction intervals.

  Examples:
      Forecast daily sales for the next 7 days based on historical data from
      a BigQuery table:

          >>> forecast(
          ...     project_id="my-gcp-project",
          ...     history_data="my-dataset.my-sales-table",
          ...     timestamp_col="sale_date",
          ...     data_col="daily_sales",
          ...     horizon=7
          ... )
          {
            "status": "SUCCESS",
            "rows": [
              {
                "forecast_timestamp": "2025-01-08T00:00:00",
                "forecast_value": 12345.67,
                "confidence_level": 0.95,
                "prediction_interval_lower_bound": 11000.0,
                "prediction_interval_upper_bound": 13691.34,
                "ai_forecast_status": ""
              },
              ...
            ]
          }

      Forecast multiple time series using a SQL query as input:

          >>> history_query = (
          ...     "SELECT unique_id, timestamp, value "
          ...     "FROM `my-project.my-dataset.my-timeseries-table` "
          ...     "WHERE timestamp > '1980-01-01'"
          ... )
          >>> forecast(
          ...     project_id="my-gcp-project",
          ...     history_data=history_query,
          ...     timestamp_col="timestamp",
          ...     data_col="value",
          ...     id_cols=["unique_id"],
          ...     horizon=14
          ... )
          {
            "status": "SUCCESS",
            "rows": [
              {
                "unique_id": "T1",
                "forecast_timestamp": "1980-08-28T00:00:00",
                "forecast_value": 1253218.75,
                "confidence_level": 0.95,
                "prediction_interval_lower_bound": 274252.51,
                "prediction_interval_upper_bound": 2232184.99,
                "ai_forecast_status": ""
              },
              ...
            ]
          }

      Error Scenarios:
          When an element in `id_cols` is not a string:

          >>> forecast(
          ...     project_id="my-gcp-project",
          ...     history_data="my-dataset.my-sales-table",
          ...     timestamp_col="sale_date",
          ...     data_col="daily_sales",
          ...     id_cols=["store_id", 123]
          ... )
          {
            "status": "ERROR",
            "error_details": "All elements in id_cols must be strings."
          }

          When `history_data` refers to a table that does not exist:

          >>> forecast(
          ...     project_id="my-gcp-project",
          ...     history_data="my-dataset.non-existent-table",
          ...     timestamp_col="sale_date",
          ...     data_col="daily_sales"
          ... )
          {
            "status": "ERROR",
            "error_details": "Not found: Table
            my-gcp-project:my-dataset.non-existent-table was not found in
            location US"
          }
  """
  model = "TimesFM 2.0"
  confidence_level = 0.95
  trimmed_upper_history_data = history_data.strip().upper()
  if trimmed_upper_history_data.startswith(
      "SELECT"
  ) or trimmed_upper_history_data.startswith("WITH"):
    history_data_source = f"({history_data})"
  else:
    history_data_source = f"TABLE `{history_data}`"

  if id_cols:
    if not all(isinstance(item, str) for item in id_cols):
      return {
          "status": "ERROR",
          "error_details": "All elements in id_cols must be strings.",
      }
    id_cols_str = "[" + ", ".join([f"'{col}'" for col in id_cols]) + "]"

    query = f"""
  SELECT * FROM AI.FORECAST(
    {history_data_source},
    data_col => '{data_col}',
    timestamp_col => '{timestamp_col}',
    model => '{model}',
    id_cols => {id_cols_str},
    horizon => {horizon},
    confidence_level => {confidence_level}
  )
  """
  else:
    query = f"""
  SELECT * FROM AI.FORECAST(
    {history_data_source},
    data_col => '{data_col}',
    timestamp_col => '{timestamp_col}',
    model => '{model}',
    horizon => {horizon},
    confidence_level => {confidence_level}
  )
  """
  return execute_sql(project_id, query, credentials, settings, tool_context)
