#### CRUD ENDPOINTS for UI Settings #####
from typing import Any, Dict, List, Union

from fastapi import APIRouter, Depends, HTTPException

import litellm
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.types.proxy.management_endpoints.ui_sso import (
    DefaultTeamSSOParams,
    SSOConfig,
)

router = APIRouter()


class IPAddress(BaseModel):
    ip: str


class SettingsResponse(BaseModel):
    """Base response model for settings with values and schema information"""

    values: Dict[str, Any]
    """The current configuration values"""

    field_schema: Dict[str, Any]
    """Schema information including descriptions and property types for UI display"""


class SSOSettingsResponse(SettingsResponse):
    """Response model for SSO settings"""

    pass


class InternalUserSettingsResponse(SettingsResponse):
    """Response model for internal user settings"""

    pass


class DefaultTeamSettingsResponse(SettingsResponse):
    """Response model for default team settings"""

    pass


@router.get(
    "/get/allowed_ips",
    tags=["Budget & Spend Tracking"],
    dependencies=[Depends(user_api_key_auth)],
    include_in_schema=False,
)
async def get_allowed_ips():
    from litellm.proxy.proxy_server import general_settings

    _allowed_ip = general_settings.get("allowed_ips")
    return {"data": _allowed_ip}


@router.post(
    "/add/allowed_ip",
    tags=["Budget & Spend Tracking"],
    dependencies=[Depends(user_api_key_auth)],
)
async def add_allowed_ip(ip_address: IPAddress):
    from litellm.proxy.proxy_server import (
        general_settings,
        prisma_client,
        proxy_config,
        store_model_in_db,
    )

    if prisma_client is None:
        raise Exception("No DB Connected")

    _allowed_ips: List = general_settings.get("allowed_ips", [])
    if ip_address.ip not in _allowed_ips:
        _allowed_ips.append(ip_address.ip)
        general_settings["allowed_ips"] = _allowed_ips
    else:
        raise HTTPException(status_code=400, detail="IP address already exists")

    if store_model_in_db is not True:
        raise HTTPException(
            status_code=500,
            detail={
                "error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature."
            },
        )

    # Load existing config
    config = await proxy_config.get_config()
    verbose_proxy_logger.debug("Loaded config: %s", config)
    if "general_settings" not in config:
        config["general_settings"] = {}

    if "allowed_ips" not in config["general_settings"]:
        config["general_settings"]["allowed_ips"] = []

    if ip_address.ip not in config["general_settings"]["allowed_ips"]:
        config["general_settings"]["allowed_ips"].append(ip_address.ip)

    await proxy_config.save_config(new_config=config)

    return {
        "message": f"IP {ip_address.ip} address added successfully",
        "status": "success",
    }


@router.post(
    "/delete/allowed_ip",
    tags=["Budget & Spend Tracking"],
    dependencies=[Depends(user_api_key_auth)],
)
async def delete_allowed_ip(ip_address: IPAddress):
    from litellm.proxy.proxy_server import general_settings, proxy_config

    _allowed_ips: List = general_settings.get("allowed_ips", [])
    if ip_address.ip in _allowed_ips:
        _allowed_ips.remove(ip_address.ip)
        general_settings["allowed_ips"] = _allowed_ips
    else:
        raise HTTPException(status_code=404, detail="IP address not found")

    # Load existing config
    config = await proxy_config.get_config()
    verbose_proxy_logger.debug("Loaded config: %s", config)
    if "general_settings" not in config:
        config["general_settings"] = {}

    if "allowed_ips" not in config["general_settings"]:
        config["general_settings"]["allowed_ips"] = []

    if ip_address.ip in config["general_settings"]["allowed_ips"]:
        config["general_settings"]["allowed_ips"].remove(ip_address.ip)

    await proxy_config.save_config(new_config=config)

    return {"message": f"IP {ip_address.ip} deleted successfully", "status": "success"}


async def _get_settings_with_schema(
    settings_key: str,
    settings_class: Any,
    config: dict,
) -> dict:
    """
    Common utility function to get settings with schema information.

    Args:
        settings_key: The key in litellm_settings to get
        settings_class: The Pydantic class to use for schema
        config: The config dictionary
    """
    from pydantic import TypeAdapter

    litellm_settings = config.get("litellm_settings", {}) or {}
    settings_data = litellm_settings.get(settings_key, {}) or {}

    # Create the settings object
    settings = settings_class(**(settings_data))
    # Get the schema
    schema = TypeAdapter(settings_class).json_schema(by_alias=True)

    # Convert to dict for response
    settings_dict = settings.model_dump()

    # Add descriptions to the response
    result = {
        "values": settings_dict,
        "field_schema": {
            "description": schema.get("description", ""),
            "properties": {},
        },
    }

    # Add property descriptions
    for field_name, field_info in schema["properties"].items():
        result["field_schema"]["properties"][field_name] = {
            "description": field_info.get("description", ""),
            "type": field_info.get("type", "string"),
        }

    # Add nested object descriptions
    for def_name, def_schema in schema.get("definitions", {}).items():
        result["field_schema"][def_name] = {
            "description": def_schema.get("description", ""),
            "properties": {
                prop_name: {"description": prop_info.get("description", "")}
                for prop_name, prop_info in def_schema.get("properties", {}).items()
            },
        }

    return result


@router.get(
    "/get/internal_user_settings",
    tags=["SSO Settings"],
    dependencies=[Depends(user_api_key_auth)],
    response_model=InternalUserSettingsResponse,
)
async def get_internal_user_settings():
    """
    Get all SSO settings from the litellm_settings configuration.
    Returns a structured object with values and descriptions for UI display.
    """
    from litellm.proxy.proxy_server import proxy_config

    # Load existing config
    config = await proxy_config.get_config()

    return await _get_settings_with_schema(
        settings_key="default_internal_user_params",
        settings_class=DefaultInternalUserParams,
        config=config,
    )


@router.get(
    "/get/default_team_settings",
    tags=["SSO Settings"],
    dependencies=[Depends(user_api_key_auth)],
    response_model=DefaultTeamSettingsResponse,
)
async def get_default_team_settings():
    """
    Get all SSO settings from the litellm_settings configuration.
    Returns a structured object with values and descriptions for UI display.
    """
    from litellm.proxy.proxy_server import proxy_config

    # Load existing config
    config = await proxy_config.get_config()

    return await _get_settings_with_schema(
        settings_key="default_team_params",
        settings_class=DefaultTeamSSOParams,
        config=config,
    )


async def update_default_team_member_budget(
    teams: List[NewUserRequestTeam], user_api_key_dict: UserAPIKeyAuth
):
    """
    1. Update the max member budget for the team
    """
    from fastapi import Request

    from litellm.proxy.management_endpoints.team_endpoints import update_team

    for team in teams:
        team_id = team.team_id
        max_budget_in_team = team.max_budget_in_team
        await update_team(
            data=UpdateTeamRequest(
                team_id=team_id,
                team_member_budget=max_budget_in_team,
            ),
            user_api_key_dict=user_api_key_dict,
            http_request=Request(scope={"type": "http"}),
        )


async def _update_litellm_setting(
    settings: Union[DefaultInternalUserParams, DefaultTeamSSOParams],
    settings_key: str,
    in_memory_var: Any,
    success_message: str,
):
    """
    Common utility function to update `litellm_settings` in both memory and config.

    Args:
        settings: The settings object to update
        settings_key: The key in litellm_settings to update
        in_memory_var: The in-memory variable to update
        success_message: Message to return on success
    """
    from litellm.proxy.proxy_server import proxy_config, store_model_in_db

    if store_model_in_db is not True:
        raise HTTPException(
            status_code=500,
            detail={
                "error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature."
            },
        )

    # Update the in-memory settings
    in_memory_var = settings.model_dump(exclude_none=True)

    # Load existing config
    config = await proxy_config.get_config()

    # Update config with new settings
    if "litellm_settings" not in config:
        config["litellm_settings"] = {}

    config["litellm_settings"][settings_key] = settings.model_dump(exclude_none=True)

    # Save the updated config
    await proxy_config.save_config(new_config=config)

    return {
        "message": success_message,
        "status": "success",
        "settings": in_memory_var,
    }


@router.patch(
    "/update/internal_user_settings",
    tags=["SSO Settings"],
    dependencies=[Depends(user_api_key_auth)],
)
async def update_internal_user_settings(
    settings: DefaultInternalUserParams,
    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
    """
    Update the default internal user parameters for SSO users.
    These settings will be applied to new users who sign in via SSO.
    """
    if settings.teams is not None and all(
        isinstance(team, NewUserRequestTeam) for team in settings.teams
    ):
        await update_default_team_member_budget(
            settings.teams, user_api_key_dict=user_api_key_dict  # type: ignore
        )

    return await _update_litellm_setting(
        settings=settings,
        settings_key="default_internal_user_params",
        in_memory_var=litellm.default_internal_user_params,
        success_message="Internal user settings updated successfully",
    )


@router.patch(
    "/update/default_team_settings",
    tags=["SSO Settings"],
    dependencies=[Depends(user_api_key_auth)],
)
async def update_default_team_settings(settings: DefaultTeamSSOParams):
    """
    Update the default team parameters for SSO users.
    These settings will be applied to new teams created from SSO.
    """
    return await _update_litellm_setting(
        settings=settings,
        settings_key="default_team_params",
        in_memory_var=litellm.default_team_params,
        success_message="Default team settings updated successfully",
    )


@router.get(
    "/get/sso_settings",
    tags=["SSO Settings"],
    dependencies=[Depends(user_api_key_auth)],
    response_model=SSOSettingsResponse,
)
async def get_sso_settings():
    """
    Get all SSO configuration settings from the environment variables.
    Returns a structured object with values and descriptions for UI display.
    """
    import os

    from litellm.proxy.proxy_server import proxy_config

    # Load existing config to get both environment variables and general settings
    config = await proxy_config.get_config()
    general_settings = config.get("general_settings", {}) or {}
    environment_variables = config.get("environment_variables", {}) or {}

    # Get user_email from general_settings
    proxy_admin_email = general_settings.get("proxy_admin_email", None)

    # Helper function to get env var value (first from config, then from environment)
    def get_env_value(env_var_name: str):
        return environment_variables.get(env_var_name) or os.getenv(env_var_name)

    # Get current environment variables for SSO
    sso_config = SSOConfig(
        google_client_id=get_env_value("GOOGLE_CLIENT_ID"),
        google_client_secret=get_env_value("GOOGLE_CLIENT_SECRET"),
        microsoft_client_id=get_env_value("MICROSOFT_CLIENT_ID"),
        microsoft_client_secret=get_env_value("MICROSOFT_CLIENT_SECRET"),
        microsoft_tenant=get_env_value("MICROSOFT_TENANT"),
        generic_client_id=get_env_value("GENERIC_CLIENT_ID"),
        generic_client_secret=get_env_value("GENERIC_CLIENT_SECRET"),
        generic_authorization_endpoint=get_env_value("GENERIC_AUTHORIZATION_ENDPOINT"),
        generic_token_endpoint=get_env_value("GENERIC_TOKEN_ENDPOINT"),
        generic_userinfo_endpoint=get_env_value("GENERIC_USERINFO_ENDPOINT"),
        proxy_base_url=get_env_value("PROXY_BASE_URL"),
        user_email=proxy_admin_email,  # Get from config instead of environment
        ui_access_mode=general_settings.get("ui_access_mode", None),
    )

    # Get the schema for UI display
    from pydantic import TypeAdapter

    schema = TypeAdapter(SSOConfig).json_schema(by_alias=True)

    # Convert to dict for response
    sso_dict = sso_config.model_dump()

    # Add descriptions to the response
    result = {
        "values": sso_dict,
        "field_schema": {
            "description": schema.get("description", ""),
            "properties": {},
        },
    }

    # Add property descriptions
    for field_name, field_info in schema["properties"].items():
        result["field_schema"]["properties"][field_name] = {
            "description": field_info.get("description", ""),
            "type": field_info.get("type", "string"),
        }

    return result


@router.patch(
    "/update/sso_settings",
    tags=["SSO Settings"],
    dependencies=[Depends(user_api_key_auth)],
)
async def update_sso_settings(sso_config: SSOConfig):
    """
    Update SSO configuration by saving to both environment variables and config file.
    """
    import os

    from litellm.proxy.proxy_server import proxy_config

    # Update environment variables
    env_var_mapping = {
        "google_client_id": "GOOGLE_CLIENT_ID",
        "google_client_secret": "GOOGLE_CLIENT_SECRET",
        "microsoft_client_id": "MICROSOFT_CLIENT_ID",
        "microsoft_client_secret": "MICROSOFT_CLIENT_SECRET",
        "microsoft_tenant": "MICROSOFT_TENANT",
        "generic_client_id": "GENERIC_CLIENT_ID",
        "generic_client_secret": "GENERIC_CLIENT_SECRET",
        "generic_authorization_endpoint": "GENERIC_AUTHORIZATION_ENDPOINT",
        "generic_token_endpoint": "GENERIC_TOKEN_ENDPOINT",
        "generic_userinfo_endpoint": "GENERIC_USERINFO_ENDPOINT",
        "proxy_base_url": "PROXY_BASE_URL",
    }

    # Load existing config
    config = await proxy_config.get_config()

    # Update config with new environment variables
    if "environment_variables" not in config:
        config["environment_variables"] = {}

    # Update general_settings for user_email (admin email)
    if "general_settings" not in config:
        config["general_settings"] = {}

    # Update environment variables in config and in memory
    sso_data = sso_config.model_dump(exclude_none=True)
    for field_name, value in sso_data.items():

        if field_name == "user_email" and value is not None:
            # Store user_email in general_settings instead of environment variables
            config["general_settings"]["proxy_admin_email"] = value
        elif field_name == "ui_access_mode" and value is not None:

            config["general_settings"]["ui_access_mode"] = value
        elif field_name in env_var_mapping and value is not None:
            env_var_name = env_var_mapping[field_name]
            # Update in config
            config["environment_variables"][env_var_name] = value
            # Update in runtime environment
            os.environ[env_var_name] = value

    stored_config = config
    if len(config["environment_variables"]) > 0:

        stored_config["environment_variables"] = proxy_config._encrypt_env_variables(
            environment_variables=config["environment_variables"]
        )
    # Save the updated config
    await proxy_config.save_config(new_config=stored_config)

    return {
        "message": "SSO settings updated successfully",
        "status": "success",
        "settings": sso_data,
    }