"""Lambda Functions for observability_service_level_prometheus_metrics_upload_delay_histogram"""

import json
import os
from datetime import datetime, timezone
from typing import Any

import boto3
import duckdb
import pymysql
import sliutil.connection
import sliutil.s3
import sliutil.tidb

SLI_TABLE_NAME = "sli_o11y_dedicated_prometheus_upload_delay"
SLI_METRIC_NAME = (
    "observability_service_level_prometheus_metrics_upload_delay_histogram"
)
SLI_S3_BUCKET_NAME = os.environ["SLI_S3_BUCKET_NAME"]


def lambda_handler(event, context):  # pylint:disable=unused-argument
    """AWS Lambda Function Handler
    Parameters
    ----------
    event: dict, required
        API Gateway Lambda Proxy Input Format
        Event doc: https://docs.aws.amazon.com/apigateway/latest/developerguide/set-up-lambda-proxy-integrations.html#api-gateway-simple-proxy-for-lambda-input-format # pylint:disable=line-too-long # noqa
    context: object, required
        Lambda Context runtime methods and attributes
        Context doc: https://docs.aws.amazon.com/lambda/latest/dg/python-context-object.html
    Returns
    ------
    API Gateway Lambda Proxy Output Format: dict
        Return doc: https://docs.aws.amazon.com/apigateway/latest/developerguide/set-up-lambda-proxy-integrations.html # pylint:disable=line-too-long
    """

    tidb_conn = sliutil.connection.get_tidb_connection()
    duckdb_conn = sliutil.connection.get_duckdb_connection()

    last_ts = float(
        sliutil.tidb.find_last_timestamp(connection=tidb_conn, table=SLI_TABLE_NAME)
    )
    begin_ts = event.get("begin_ts", last_ts - 10 * 60)
    end_ts = event.get("end_ts", last_ts + 1 * 60 * 60)
    end_ts = min(end_ts, datetime.now(timezone.utc).timestamp())

    handler = Handler(tidb_conn=tidb_conn, duckdb_conn=duckdb_conn)
    handler.run(begin_ts=begin_ts, end_ts=end_ts)

    return {
        "statusCode": 200,
        "body": json.dumps(
            {
                "message": "hello world",
                # "location": ip.text.replace("\n", "")
            }
        ),
    }


class Handler:
    """Handle metrics `observability_service_level_prometheus_metrics_upload_delay_histogram`"""

    def __init__(
        self, tidb_conn: pymysql.Connection, duckdb_conn: duckdb.DuckDBPyConnection
    ) -> None:
        self.data = []
        self.probe_counter = {}
        self.tidb_conn = tidb_conn
        self.duckdb_conn = duckdb_conn
        self.s3_client = boto3.client("s3")

    def run(self, begin_ts: float, end_ts: float) -> None:
        """Extract metrics from the original table from begin_ts to end_ts
        and then save them into a new table
        """
        print(f"Start, begin_ts: {begin_ts}, end_ts: {end_ts}")
        rows = self.search(begin_ts=begin_ts, end_ts=end_ts)
        for row in rows:
            metric_timestamp = row[0]
            instance = row[1]
            metric_group = json.loads(row[2])
            metrics = metric_group["metrics"]

            for metric in metrics:
                buckets = metric["buckets"]
                for key, value in buckets.items():
                    tmp = self.parse(
                        timestamp=metric_timestamp,
                        instance=instance,
                        metric=metric,
                        bucket_key=key,
                        bucket_value=value,
                    )
                    if tmp is not None:
                        self.data.append(tmp)
                    if len(self.data) >= 1000:
                        self.save()
                        self.data = []

        self.save()
        print("Finish")

    def parse(
        self,
        timestamp: int,
        instance: str,
        metric: dict[str, Any],
        bucket_key: str,
        bucket_value: str,
    ) -> tuple | None:
        """Parse the input metric, return a tuple about probe info"""
        key = (
            f"{metric['labels']['o11y_cluster_id']}-{instance}"
            f"-{metric['labels']['status']}-{bucket_key}"
        )
        current_value = int(bucket_value)
        if key not in self.probe_counter:
            self.probe_counter[key] = current_value
            return None
        increment = 0
        if current_value >= self.probe_counter[key]:
            increment = current_value - self.probe_counter[key]
        self.probe_counter[key] = current_value

        return (
            timestamp,
            metric["labels"]["o11y_cluster_id"],
            instance,
            metric["labels"]["region"],
            metric["labels"]["tenant_id"],
            metric["labels"]["project_id"],
            metric["labels"]["application_id"],
            metric["labels"]["status"],
            bucket_key,
            current_value,
            increment,
        )

    def search(self, begin_ts: float, end_ts: float) -> list:
        """List all metrics from begin_ts to end_ts"""
        file_names = sliutil.s3.ls_s3_objects(
            self.s3_client,
            begin_ts=begin_ts,
            end_ts=end_ts,
            bucket_name=SLI_S3_BUCKET_NAME,
            metric_name=SLI_METRIC_NAME,
        )
        if len(file_names) == 0:
            return []
        start = datetime.now(timezone.utc).timestamp()
        sql = f"""
            SELECT timestamp, instance, metrics
            FROM read_parquet({file_names}, hive_partitioning=false)
        """
        # print(f"search, sql: {sql}")
        result = self.duckdb_conn.execute(sql).fetchall()
        end = datetime.now(timezone.utc).timestamp()
        print(f"search, take {end - start:.2f}s")
        return result

    def save(self) -> None:
        """Save data into table"""
        # print(f"save data, len: {len(self.data)}")
        with self.tidb_conn.cursor() as cursor:
            sql = f"""
                REPLACE INTO `{SLI_TABLE_NAME}` (
                    timestamp, o11y_cluster_id, instance, region, tenant_id,
                    project_id, application_id, status, le, value, increment
                )
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            """
            cursor.executemany(sql, self.data)
            self.tidb_conn.commit()