# ruff: noqa: PLW0603
import logging
import os
import sys
import uuid

from alphatrion import envs
from alphatrion.storage import runtime as storage_runtime
from alphatrion.storage.sqlstore import SQLStore
from alphatrion.utils.pricing import load_pricing_config

logger = logging.getLogger(__name__)

__RUNTIME__ = None


def init(
    user_id: uuid.UUID,
    team_id: uuid.UUID | None = None,
    org_id: uuid.UUID | None = None,
):
    """
    Initialize the AlphaTrion runtime environment.

    Args:
        user_id: The user ID for the current user. You can generate a UUID
                 using `uuid.uuid4()`.
        team_id: The team ID for the current user. If not provided, will look
                 for the first team associated with the user in the database.
    """
    global __RUNTIME__
    __RUNTIME__ = Runtime(
        user_id=user_id,
        team_id=team_id,
        org_id=org_id,
    )


def global_runtime():
    if __RUNTIME__ is None:
        raise RuntimeError("Runtime is not initialized. Call alphatrion.init() first.")
    return __RUNTIME__


# Runtime contains all kinds of clients, e.g., metadb client, artifact client, etc.
# Stateful information will also be stored here.
class Runtime:
    __slots__ = (
        "_user_id",
        "_team_id",
        "_org_id",
        "_metadb",
        "_tracestore",
        "_artifact",
        "_root_path",
        "_current_experiment",
    )

    def __init__(
        self,
        user_id: uuid.UUID,
        team_id: uuid.UUID | None = None,
        org_id: uuid.UUID | None = None,
    ):
        # Load pricing config at startup - exit if it fails
        try:
            load_pricing_config()
            logger.info("Successfully loaded model pricing configuration")
        except Exception as e:
            logger.error(f"Failed to load pricing configuration: {e}")
            logger.error("Application cannot start without valid pricing configuration")
            sys.exit(1)

        storage_runtime.init()
        self._metadb = storage_runtime.storage_runtime().metadb
        self._tracestore = storage_runtime.storage_runtime().tracestore
        self._artifact = storage_runtime.storage_runtime().artifact

        self._user_id = user_id
        self._team_id = team_id
        self._org_id = org_id

        if team_id is None:
            # If team_id is not provided, look for the first team associated with
            # the user in the database.
            teams = self._metadb.list_user_teams(user_id)
            if len(teams) == 0:
                raise ValueError(
                    f"No team found for user_id {user_id}. Make sure the user is "
                    f"associated with at least one team in the database."
                )
            self._team_id = teams[0].uuid

        if org_id is None:
            # Look up org_id from user or team
            user = self._metadb.get_user(user_id)
            if not user:
                raise ValueError(f"User {user_id} not found in the database.")
            self._org_id = user.org_id

        self._root_path = os.getenv(envs.ROOT_PATH, os.path.expanduser("~/.alphatrion"))
        if not os.path.exists(self._root_path):
            os.makedirs(self._root_path, exist_ok=True)

    @property
    def metadb(self) -> SQLStore:
        return self._metadb

    @property
    def tracestore(self):
        return self._tracestore

    @property
    def artifact(self):
        return self._artifact

    @property
    def user_id(self) -> uuid.UUID:
        return self._user_id

    @property
    def team_id(self) -> uuid.UUID:
        return self._team_id

    @property
    def org_id(self) -> uuid.UUID:
        return self._org_id

    @property
    def root_path(self) -> str:
        return self._root_path

    @property
    def current_experiment(self):
        return self._current_experiment

    @current_experiment.setter
    def current_experiment(self, exp):
        self._current_experiment = exp
