import uuid
import time
import json

import boto3
from prometheus_client import Counter, Histogram
from flask import request

from consts.consts import EDA_PROMETHEUS_BUCKETS, ENV_DEV, ENV_STAGING, ENV_PROD, MEASURE_SKU_NETWORK, MEASURE_SKU_REQUESTS, MEASURE_SKU_TOKENS
from config import config
from errors.errors import BadRequestException
from handlers.base import get_access_key, get_app_id, get_endpoint, get_org_id_from_header

from models.base import rds
from utils.local import get_local_metrics
from utils.logging import logging


EDA_OPENAI_HISTOGRAM = Histogram("eda_openai_histogram", "openai histogram", labelnames=("orgId", "appId", "endpoint", "appKey", "status", "attempt"), buckets=EDA_PROMETHEUS_BUCKETS)
EDA_TOKEN_COUNTER = Counter("eda_token_counter", "token counter", labelnames=("orgId", "appId", "endpoint", "appKey"))
EDA_REQUESTS_COUNTER = Counter("eda_request_counter", "eda requests", labelnames=("orgId", "appId", "endpoint", "appKey"))
EDA_NETWORK_COUNTER = Counter("eda_network_counter", "network usage counter", labelnames=("orgId", "appId", "endpoint", "appKey"))
EDA_DATAAPI_HISTOGRAM = Histogram("eda_dataapi_histogram", "dataapi histogram", labelnames=("region", "url", "status"), buckets=EDA_PROMETHEUS_BUCKETS)


sqs_queue_url_mapper = {
    ENV_DEV: "https://sqs.us-west-2.amazonaws.com/385595570414/billing-report-nightly",
    ENV_STAGING: "https://sqs.us-west-2.amazonaws.com/385595570414/billing-report-staging",
    ENV_PROD: "https://sqs.us-west-2.amazonaws.com/380838443567/billing-report-prod",
}


class BillingMeasure:
    def __init__(self, tenant_id: int, cost_id: str, cost: int, cost_timestamp: int, labels: dict[str, str], unit: str):
        self.tenant_id = tenant_id  # 租户ID
        self.cost_id = cost_id  # 唯一ID例如UUID
        self.cost = cost  # 数量
        self.cost_timestamp = cost_timestamp  # 时间戳，单位秒
        self.labels = labels
        self.unit = unit  # sku

    def to_json(self):
        return {
            "tenant_id": self.tenant_id,
            "cost_id": self.cost_id,
            "cost": self.cost,
            "cost_time_stamp": self.cost_timestamp,
            "labels": self.labels,
            "unit": self.unit,
        }


def hold_measure_metric(measure: BillingMeasure):
    if "appId" not in measure.labels:
        raise BadRequestException("invalid measure metrics, appId is required")

    labels = {}
    for k, v in measure.labels.items():
        labels[str(k)] = str(v)

    org_id = measure.tenant_id
    measure.labels = labels
    key = f"measure:{org_id}"
    rds.lpush(key, json.dumps(measure.to_json()))
    logging.info("hold_measure_metric %s", measure)


def submit_measure_metrics(measure: dict):
    logging.info("start to send %s to billing measure sqs", measure)

    sqs = boto3.client('sqs')
    queue_url = sqs_queue_url_mapper[config.sqs_region_env]
    message_body = json.dumps(measure)

    # Send message to SQS queue
    response = sqs.send_message(
        QueueUrl=queue_url,
        DelaySeconds=0,
        MessageBody=message_body,
    )

    logging.info("send %s to billing measure sqs got : %s", message_body, response['MessageId'])


def measuring_openai(stage, model, response):
    metrics = get_local_metrics()

    if not metrics:
        logging.info("no metrics found in measuring_openai, ignore")
        return

    if response is None:
        logging.info("response is empty, reject measuring")
        return

    org_id = metrics["org_id"]
    cost_id = "{}-{}".format(metrics["job_id"], int(time.time() * 1000))
    app_id = metrics["app_id"]
    endpoint = metrics["endpoint"]
    app_key = metrics["app_key"]
    total_tokens = response.usage.total_tokens
    EDA_TOKEN_COUNTER.labels(org_id, app_id, endpoint, app_key).inc(total_tokens)
    logging.info(
        "measuring_openai, org_id, app_id, app_key, endpoint, total_tokens: %s, %s, %s, %s, %s",
        org_id, app_id, app_key, endpoint, total_tokens,
    )

    labels = {
        "data_context_id": metrics["data_context_id"],
        "session_context_id": metrics["session_context_id"],
        "stage": stage,
        "model": model,
        "appId": app_id,
    }
    measure = BillingMeasure(int(org_id), cost_id, total_tokens, int(time.time()), labels, MEASURE_SKU_TOKENS)
    try:
        hold_measure_metric(measure)
    except Exception as e:
        logging.error("failed to submit measure metrics: %s", e)


def measuring_api():
    if config.DISABLE_MEASURING:
        logging.info("measuring is disabled")
        return

    org_id = get_org_id_from_header()
    assert org_id is not None
    cost_id = "{}-{}".format(request.headers.get("x-debug-trace-id", uuid.uuid4().hex), int(time.time() * 100))
    app_id = get_app_id()
    app_key = get_access_key()
    endpoint = get_endpoint()

    logging.info("measuring_api, org_id, app_id, app_key, endpoint: %s, %s, %s, %s", org_id, app_id, app_key, endpoint)
    EDA_REQUESTS_COUNTER.labels(org_id, app_id, endpoint, app_key).inc()

    labels = {
        "endpoint": endpoint,
        "app_key": app_key,
        "appId": app_id,
    }
    measure = BillingMeasure(int(org_id), cost_id, 1, int(time.time()), labels, MEASURE_SKU_REQUESTS)
    try:
        hold_measure_metric(measure)
    except Exception as e:
        logging.error("failed to hold measure metrics: %s", e)


def measuring_network(org_id: int, app_id: int, app_key: str, bytes_count: int):
    logging.info("measure_network, org_id, app_id, bytes_count: %s, %s, %s", org_id, app_id, bytes_count)
    endpoint = "POST - /v2/chat2data"
    EDA_NETWORK_COUNTER.labels(org_id, app_id, endpoint, app_key).inc(bytes_count)

    cost_id = uuid.uuid4().hex

    labels = {
        "appId": str(app_id),
    }
    measure = BillingMeasure(int(org_id), cost_id, bytes_count, int(time.time()), labels, MEASURE_SKU_NETWORK)
    try:
        hold_measure_metric(measure)
    except Exception as e:
        logging.error("failed to hold measure metrics: %s", e)