from re import sub
from traceback import format_exc
from datetime import datetime
from http import HTTPStatus
from logging import Logger, getLogger
from os import environ
from warnings import filterwarnings
from typing import Dict
from pydantic import BaseModel, ConfigDict, ValidationError

# import asyncio

from prometheus_client import Counter

import bentoml
from bentoml.exceptions import BadInput, InternalServerError
from pandas import DataFrame
from mlflow.exceptions import MlflowException

from internal_lib.serving import Transformer
from internal_lib.serving.extractors import (
    Extractor1,
    Extractor2,
    Extractor3,
    Extractor4,
    Extractor5,
    Extractor6
)

API_PORT: int = int(environ.get("API_PORT", 5000))
API_TIMEOUT: int = int(environ.get("API_TIMEOUT", 60))
API_MAX_CONCURRENCY: int = int(environ.get("API_MAX_CONCURRENCY", 50))
API_LOGLEVEL: str = environ.get("API_LOGLEVEL", "INFO")
API_WORKERS: int | str = environ.get("API_WORKERS", "cpu_count")
API_WORKERS = int(API_WORKERS) if API_WORKERS.isdigit() else API_WORKERS
# ENABLE_ASYNCIO: str = environ.get("ENABLE_ASYNCIO", "False")

MLFLOW_MODEL_NAME: str = environ.get("MLFLOW_MODEL_NAME", "NO_VALUE_DEFINED")
MLFLOW_MODEL_VERSION: str = environ.get("MLFLOW_MODEL_VERSION", "NO_VALUE_DEFINED")

MODEL_NAME: str = environ.get("MODEL_NAME", "NO_VALUE_DEFINED")
MODEL_VERSION: str = environ.get("MODEL_VERSION", "NO_VALUE_DEFINED")
API_SERVICE_NAME: str = sub(r"[-.]", "_", f"API_SERVICE_{MODEL_NAME}").upper()
API_PROFILLING_ENABLE: str = environ.get("API_PROFILLING_ENABLE", "")

if API_PROFILLING_ENABLE == "True":
    from pyinstrument import Profiler
    from pyinstrument.renderers import ConsoleRenderer


ENABLE_LOGGING_ACCESS: bool = API_LOGLEVEL == "DEBUG"

bentoml_logger: Logger = getLogger("api_model")
bentoml_logger.setLevel(API_LOGLEVEL)

if API_LOGLEVEL == "DEBUG":
    bentoml_logger.debug(
        msg=f"""
API_SERVICE_NAME: {API_SERVICE_NAME}
API_PORT: {API_PORT}
API_WORKERS: {API_WORKERS}
API_WORKER_TIMEOUT: {API_TIMEOUT}
API_WORKER_MAX_CONCURRENCY: {API_MAX_CONCURRENCY}
API_LOGLEVEL: {API_LOGLEVEL}
MODEL_NAME: {MODEL_NAME}
MODEL_VERSION: {MODEL_VERSION}"""
    )
else:
    filterwarnings("ignore")
    getLogger("bentoml.access").disabled = True
    getLogger("bentoml").disabled = True
    getLogger(__name__).disabled = True
    getLogger("bentoml.serve").disabled = True
    getLogger("bentoml._internal.tag").disabled = True


def return_error_message(predict_id: int | str, exception: str) -> dict:
    error: dict = {
        "predict_id": predict_id,
        "error": exception,
        "model_name": MODEL_NAME,
        "model_version": MODEL_VERSION,
        "traceback": sub(r'[\\^"\n]|  ', "", format_exc()),
    }
    return error

class DataModel(BaseModel):
    model_config = ConfigDict(extra="allow")
    data_1: Dict
    data_2: Dict | None = None
    data_3: Dict
    data_4: Dict

api_worker_counter = Counter(
    name='bentoml_service_api_workers_total',
    documentation='Total number of API Workers'
)
api_worker_counter.inc(1)

@bentoml.service(
    name=API_SERVICE_NAME,
    workers=API_WORKERS,
    http={"port": API_PORT},
    traffic={
        "timeout": API_TIMEOUT,
        "max_concurrency": API_MAX_CONCURRENCY,
    },
    metrics={
        "enabled": True,
        "duration": {"min": 0.05, "max": 5.0, "factor": 1.5}
    },
    logging={
        "access": {
            "enabled": ENABLE_LOGGING_ACCESS,
            "request_content_length": True,
            "request_content_type": True,
            "response_content_length": True,
            "response_content_type": True,
            "skip_paths": ["/metrics", "/healthz", "/livez", "/readyz"],
            "format": {"trace_id": "032x", "span_id": "016x"},
        }
    },
)
class APIService:
    bentoml_logger.info("Loading model from MLFlow: %s - %s", MLFLOW_MODEL_NAME, MLFLOW_MODEL_VERSION)

    def __init__(self) -> None:

        self.worker_index: int = bentoml.server_context.worker_index
        self.bento_model = bentoml.mlflow.load_model("model:latest")
        if API_PROFILLING_ENABLE == "True":
            self.counter = 0
            self.profiler = Profiler(interval=0.01, use_timing_thread=True)

        self.transform = Transformer(
            [
                Extractor1(),
                Extractor2(),
                Extractor3(),
                Extractor4(),
                Extractor5(),
                Extractor6()
            ],
        )

        bentoml_logger.info("Worker %s ready.", self.worker_index)

    @bentoml.api(input_spec=dict, route="/api/")
    # async def predict(self, **params: dict):
    def predict(self, **params: dict):
        try:
            if API_PROFILLING_ENABLE == "True":
                self.profiler.start()
            payload: dict = params.get("root")
            predict_id: int | str = payload.get("predict_id", "not found")

            predictDataModel(**(params.get("root")))
            # if ENABLE_ASYNCIO == "True":
            #     features: DataFrame = await asyncio.to_thread(self.transform.run, payload)
            # else:
            features: DataFrame = self.transform.run(payload)

            # if ENABLE_ASYNCIO == "True":
            #     prediction = await asyncio.to_thread(self.bento_model.predict, features)
            #     score: float = prediction[0, 1]
            # else:
            score: float = self.bento_model.predict(features)[0, 1]

            output: dict = {
                "score": score,
                "features": features.fillna("None").astype(str).to_dict("records")[0],
                "datetime": f"{str(datetime.now())}, 'info'",
            }

            response: dict = {
                "typeName": "...",
                "name": MODEL_NAME,
                "value": score,
                "details": output,
                "version": MODEL_VERSION,
                "state": 0,
            }

            bentoml_logger.info(
                "predict ID: %s - Model: %s - Version %s - Score: %s",
                predict_id,
                MODEL_NAME,
                MODEL_VERSION,
                round(score, 3)
            )
            if API_PROFILLING_ENABLE == "True":
                self.counter += 1
                if self.counter % 100 == 0:
                    sessionz = self.profiler.stop()
                    profile_renderer = ConsoleRenderer(
                        unicode=False,
                        color=True,
                        flat=False,
                        time='seconds',
                        flat_time='self',
                        short_mode=True,
                        show_all=False,
                        timeline=False,
                    )
                    print(profile_renderer.render(sessionz))
            return response
        except (
            AttributeError,
            TypeError,
            IndexError,
            ValueError,
            KeyError,
            AssertionError,
            ValidationError,
            MlflowException,
        ) as exception:
            raise BadInput(
                message=return_error_message(predict_id=predict_id, exception=str(exception))
            )
        except Exception as exception:
            raise InternalServerError(
                message=return_error_message(predict_id=predict_id, exception=str(exception)),
                error_code=HTTPStatus.INTERNAL_SERVER_ERROR,
            )