"""
Config repository for database operations on LiteLLM_Config.

This repository handles config reconciliation between database values and
YAML configmap values. DB values override configmap values except for
None values and empty lists.
"""

import asyncio
import copy
import json
import os
from typing import Any, Dict, List, Literal, Optional, cast

from litellm._logging import verbose_proxy_logger
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper


class ConfigParam:
    """Simple wrapper for config parameter from DB."""

    def __init__(self, param_name: str, param_value: Any):
        self.param_name = param_name
        self.param_value = param_value


class ConfigRepository:
    """Repository for config database operations with reconciliation support."""

    CONFIG_PARAMS = [
        "general_settings",
        "router_settings",
        "litellm_settings",
        "environment_variables",
    ]

    def __init__(self, prisma_client: Any):
        self._prisma_client = prisma_client

    @property
    def prisma_client(self) -> Any:
        if self._prisma_client is None:
            raise RuntimeError(
                "No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys"
            )
        return self._prisma_client

    @property
    def table(self) -> Any:
        return self.prisma_client.db.litellm_config

    async def get_param(self, param_name: str) -> Optional[ConfigParam]:
        """Get a config parameter from the database."""
        record = await self.table.find_unique(where={"param_name": param_name})
        if record is None:
            return None
        param_value = record.param_value
        if isinstance(param_value, str):
            param_value = json.loads(param_value)
        return ConfigParam(param_name=param_name, param_value=param_value)

    async def set_param(self, param_name: str, param_value: Any) -> ConfigParam:
        """Set a config parameter in the database."""
        value_json = (
            json.dumps(param_value) if not isinstance(param_value, str) else param_value
        )
        await self.table.upsert(
            where={"param_name": param_name},
            data={
                "create": {"param_name": param_name, "param_value": value_json},
                "update": {"param_value": value_json},
            },
        )
        return ConfigParam(param_name=param_name, param_value=param_value)

    async def delete_param(self, param_name: str) -> bool:
        """Delete a config parameter from the database."""
        try:
            await self.table.delete(where={"param_name": param_name})
            return True
        except Exception:
            return False

    async def get_all_params(self) -> Dict[str, Any]:
        """Get all config parameters from the database."""
        records = await self.table.find_many()
        result = {}
        for record in records:
            param_value = record.param_value
            if isinstance(param_value, str):
                param_value = json.loads(param_value)
            result[record.param_name] = param_value
        return result

    def _deep_merge_dicts(self, dst: dict, src: dict) -> None:
        """Deep-merge src into dst, skipping None values and empty lists from src.

        On conflicts, src (DB) wins, but empty lists are treated as "no value"
        and don't overwrite the destination.
        """
        stack = [(dst, src)]
        while stack:
            d, s = stack.pop()
            for k, v in s.items():
                if v is None:
                    continue
                if isinstance(v, list) and len(v) == 0:
                    continue
                if isinstance(v, dict) and isinstance(d.get(k), dict):
                    stack.append((d[k], v))
                else:
                    d[k] = v

    def _decrypt_env_variables(
        self, env_vars: Dict[str, Any], return_original_value: bool = True
    ) -> Dict[str, str]:
        """Decrypt environment variables from database."""
        decrypted: Dict[str, str] = {}
        for key, value in env_vars.items():
            if isinstance(value, str):
                decrypted_value = decrypt_value_helper(
                    value=value,
                    key=key,
                    exception_type="debug",
                    return_original_value=return_original_value,
                )
                if decrypted_value is not None:
                    decrypted[key] = decrypted_value
            else:
                decrypted[key] = str(value)
        return decrypted

    def _normalize_env_variable_keys(self, env_vars: Dict[str, str]) -> Dict[str, str]:
        """Normalize env variable keys to include both original and uppercase versions."""
        normalized: Dict[str, str] = {}
        for key, value in env_vars.items():
            normalized[key] = value
            upper_key = key.upper()
            normalized[upper_key] = value
        return normalized

    def _update_config_fields(
        self,
        current_config: dict,
        param_name: Literal[
            "general_settings",
            "router_settings",
            "litellm_settings",
            "environment_variables",
        ],
        db_param_value: Any,
    ) -> dict:
        """Update config fields with DB values, handling the merge strategy."""
        if param_name == "environment_variables":
            decrypted_env_vars = self._decrypt_env_variables(
                db_param_value, return_original_value=True
            )
            merged_env_vars = self._normalize_env_variable_keys(decrypted_env_vars)
            for env_key, value in merged_env_vars.items():
                os.environ[env_key] = value

            current_config.setdefault("environment_variables", {}).update(
                merged_env_vars
            )
            return current_config

        if param_name not in current_config:
            current_config[param_name] = db_param_value
            return current_config

        if isinstance(current_config[param_name], dict) and isinstance(
            db_param_value, dict
        ):
            self._deep_merge_dicts(current_config[param_name], db_param_value)
        else:
            current_config[param_name] = db_param_value

        return current_config

    async def reconcile_config(
        self,
        yaml_config: dict,
        store_model_in_db: Optional[bool] = None,
    ) -> dict:
        """Reconcile config from YAML with database overrides.

        This is the main config reconciliation method that loads config params
        from the database and merges them with the YAML config. DB values
        override YAML values except for None values and empty lists.

        Args:
            yaml_config: The configuration loaded from YAML file
            store_model_in_db: Whether to load config from DB

        Returns:
            The merged configuration with DB overrides applied
        """
        if store_model_in_db is not True:
            verbose_proxy_logger.info(
                "'store_model_in_db' is not True, skipping db config reconciliation"
            )
            return yaml_config

        tasks = [self.get_param(k) for k in self.CONFIG_PARAMS]
        responses = await asyncio.gather(*tasks)

        config = copy.deepcopy(yaml_config)
        for response in responses:
            if response is None:
                continue

            param_name = response.param_name
            param_value = response.param_value
            verbose_proxy_logger.debug(
                f"param_name={param_name}, param_value={param_value}"
            )

            if param_name is not None and param_value is not None:
                config = self._update_config_fields(
                    current_config=config,
                    param_name=cast(
                        Literal[
                            "general_settings",
                            "router_settings",
                            "litellm_settings",
                            "environment_variables",
                        ],
                        param_name,
                    ),
                    db_param_value=param_value,
                )

        return config

    async def prefetch_params(self, param_names: List[str]) -> None:
        """Prefetch config params to warm the cache.

        This can be called before reconcile_config to ensure all needed
        params are loaded in a single batch.
        """
        await asyncio.gather(*[self.get_param(k) for k in param_names])
