"""
1. Default permissions for members in a team - allowed to call /key/info and /key/health
    - Create a team, create a member in a team (role = "user")


    Invalid Permissions:
    - User tries creating a key with team_id = team_id -> expect to fail. Invalid Permissions 
    - User tries editing a key with team_id = team_id -> expect to fail. Invalid Permissions 
    - User tries deleting a key with team_id = team_id -> expect to fail. Invalid Permissions 
    - User tries regenerating a key with team_id = team_id -> expect to fail. Invalid Permissions 

    Valid Permissions:
    - User tries calling /key/info with team_id, expect to get valid response 



2. Permissions - members allowd to edit, delete keys but not allowed to create keys
    - Create a team with member_permissions = ["/key/update", "/key/delete", "/key/info"]
    - Create a member in the team with role = "user"

    Valid Permissions:
    - User tries editing a key with team_id = team_id -> expect to pass. Valid Permissions
    - Note: Delete/regenerate require key ownership or team admin status, not just team member permissions
    - User tries deleting a key with team_id = team_id -> expect to fail (403) unless user owns the key or is team admin
    - User tries regenerating a key with team_id = team_id -> expect to fail (403) unless user owns the key or is team admin

    Invalid Permissions:
    - User tries creating a key with team_id = team_id -> expect to fail. Invalid Permissions
    - User tries calling /key/info with team_id, expect to get valid response 



3. Permissions - members allowed to create keys but not allowed to edit, delete keys
    - Create a team with member_permissions = ["/key/generate"]
    - Create a member in the team with role = "user"

    Valid Permissions:
    - User tries creating a key with team_id = team_id -> expect to pass. Valid Permissions

    Invalid Permissions:
    - User tries editing a key with team_id = team_id -> expect to fail. Invalid Permissions
    - User tries deleting a key with team_id = team_id -> expect to fail. Invalid Permissions
    - User tries regenerating a key with team_id = team_id -> expect to fail. Invalid Permissions
"""

import pytest
import asyncio
import aiohttp, openai
from litellm._uuid import uuid
import json
from litellm.proxy._types import ProxyErrorTypes
from typing import Optional

LITELLM_MASTER_KEY = "sk-1234"


async def create_team(session, key, member_permissions=None):
    url = "http://0.0.0.0:4000/team/new"
    headers = {
        "Authorization": f"Bearer {key}",
        "Content-Type": "application/json",
    }
    data = {"team_member_permissions": member_permissions}

    async with session.post(url, headers=headers, json=data) as response:
        status = response.status
        response_text = await response.text()

        if status != 200:
            raise Exception(response_text)

        return await response.json()


async def create_user(session, key, user_id, team_id=None):
    url = "http://0.0.0.0:4000/user/new"
    headers = {
        "Authorization": f"Bearer {key}",
        "Content-Type": "application/json",
    }
    data = {"user_id": user_id}
    if team_id:
        data["team_id"] = team_id

    async with session.post(url, headers=headers, json=data) as response:
        status = response.status
        response_text = await response.text()

        if status != 200:
            raise Exception(response_text)

        return await response.json()


async def add_team_member(session, key, team_id, user_id, role="user"):
    url = "http://0.0.0.0:4000/team/member_add"
    headers = {
        "Authorization": f"Bearer {key}",
        "Content-Type": "application/json",
    }
    data = {"team_id": team_id, "member": {"role": role, "user_id": user_id}}
    print("Adding team member with data: ", data)
    async with session.post(url, headers=headers, json=data) as response:
        status = response.status
        response_text = await response.text()

        if status != 200:
            raise Exception(response_text)

        return await response.json()


async def generate_key(session, key, team_id=None, user_id=None):
    url = "http://0.0.0.0:4000/key/generate"
    headers = {
        "Authorization": f"Bearer {key}",
        "Content-Type": "application/json",
    }
    data = {}
    if team_id:
        data["team_id"] = team_id
    if user_id:
        data["user_id"] = user_id

    async with session.post(url, headers=headers, json=data) as response:
        status = response.status
        response_text = await response.text()

        if status != 200:
            return {"status": status, "error": response_text}

        return await response.json()


async def key_info(session, key, key_id):
    url = f"http://0.0.0.0:4000/key/info?key={key_id}"
    headers = {
        "Authorization": f"Bearer {key}",
        "Content-Type": "application/json",
    }

    async with session.get(url, headers=headers) as response:
        status = response.status
        response_text = await response.text()

        if status != 200:
            return {"status": status, "error": response_text}

        return await response.json()


async def update_key(
    session: aiohttp.ClientSession,
    key: str,
    key_id: str,
    team_id: Optional[str] = None,
):
    """
    Update a key

    Args:
        key: key to use for authentication
        key_id: key to update
    """
    url = "http://0.0.0.0:4000/key/update"
    headers = {
        "Authorization": f"Bearer {key}",
        "Content-Type": "application/json",
    }
    data = {"key": key_id, "metadata": {"updated": True}}
    if team_id:
        data["team_id"] = team_id

    async with session.post(url, headers=headers, json=data) as response:
        status = response.status
        response_text = await response.text()

        if status != 200:
            return {"status": status, "error": response_text}

        return await response.json()


async def delete_key(session, key, key_id):
    url = "http://0.0.0.0:4000/key/delete"
    headers = {
        "Authorization": f"Bearer {key}",
        "Content-Type": "application/json",
    }
    data = {"keys": [key_id]}

    async with session.post(url, headers=headers, json=data) as response:
        status = response.status
        response_text = await response.text()

        if status != 200:
            return {"status": status, "error": response_text}

        return await response.json()


async def regenerate_key(session, key, key_id, team_id=None):
    url = "http://0.0.0.0:4000/key/regenerate"
    headers = {
        "Authorization": f"Bearer {key}",
        "Content-Type": "application/json",
    }
    data = {"key": key_id}
    if team_id:
        data["team_id"] = team_id

    async with session.post(url, headers=headers, json=data) as response:
        status = response.status
        response_text = await response.text()

        if status != 200:
            return {"status": status, "error": response_text}

        return await response.json()


@pytest.mark.asyncio()
async def test_default_member_permissions():
    """
    Test default permissions for members in a team - allowed to call /key/info and /key/health
    """
    async with aiohttp.ClientSession() as session:
        master_key = LITELLM_MASTER_KEY

        # Create a team
        team_data = await create_team(session=session, key=master_key)
        team_id = team_data["team_id"]

        # create a team key
        team_key_data = await generate_key(
            session=session, key=master_key, team_id=team_id
        )
        team_key = team_key_data["key"]

        # create a user
        user_data = await create_user(
            session=session,
            key=master_key,
            user_id=f"user_{uuid.uuid4().hex[:8]}",
            team_id=team_id,
        )
        user_id = user_data["user_id"]

        # Create a user key
        print("New user data: ", user_data)

        # Create a user key
        user_key_data = await generate_key(
            session=session, key=master_key, user_id=user_id
        )
        print("new user key: ", user_key_data)
        user_key = user_key_data["key"]

        # Test invalid permissions
        # User tries creating a key with team_id
        print(
            "Regular team member trying to create a key with team_id. Expecting error."
        )
        create_result = await generate_key(
            session=session, key=user_key, team_id=team_id
        )
        print("result: ", create_result)
        assert (
            "status" in create_result and create_result["status"] == 401
        ), "User should not be able to create keys for team"
        error_data = json.loads(create_result["error"])
        print("error response =", json.dumps(error_data, indent=4))
        assert (
            error_data["error"]["type"]
            == ProxyErrorTypes.team_member_permission_error.value
        ), "Error should be a team member permission error"

        # User tries editing a key with team_id
        print("Regular team member trying to edit a key with team_id. Expecting error.")
        update_result = await update_key(
            session=session, key=user_key, key_id=team_key, team_id="ATTACKER_TEAM_ID"
        )
        assert (
            "status" in update_result and update_result["status"] == 401
        ), "User should not be able to update keys for team"
        error_data = json.loads(update_result["error"])
        print("error response =", json.dumps(error_data, indent=4))
        assert (
            error_data["error"]["type"]
            == ProxyErrorTypes.team_member_permission_error.value
        ), "Error should be a team member permission error"

        # User tries deleting a key with team_id
        print(
            "Regular team member trying to delete a key with team_id. Expecting error."
        )
        delete_result = await delete_key(
            session=session,
            key=user_key,
            key_id=team_key,
        )
        assert (
            "status" in delete_result and delete_result["status"] == 403
        ), "User should not be able to delete keys for team"
        error_data = json.loads(delete_result["error"])
        print("error response =", json.dumps(error_data, indent=4))
        # Delete endpoint now returns 403 with authorization error, not team_member_permission_error
        assert "error" in error_data, "Error should contain error field"

        # User tries regenerating a key with team_id
        print(
            "Regular team member trying to regenerate a key with team_id. Expecting error."
        )
        regenerate_result = await regenerate_key(
            session=session,
            key=user_key,
            key_id=team_key,
        )
        assert (
            "status" in regenerate_result and regenerate_result["status"] == 401
        ), "User should not be able to regenerate keys for team"
        error_data = json.loads(regenerate_result["error"])
        print("error response =", json.dumps(error_data, indent=4))
        # Regenerate endpoint now returns 403 with authorization error, not team_member_permission_error
        assert "error" in error_data, "Error should contain error field"

        # Test valid permissions
        # User tries calling /key/info with team_id
        print(
            "Regular team member trying to get key info with team_id. Expecting success."
        )
        info_result = await key_info(
            session=session,
            key=user_key,
            key_id=team_key,
        )
        print("info result =", info_result)
        assert "status" not in info_result, "Admin should be able to get key info"


@pytest.mark.asyncio()
async def test_edit_delete_permissions():
    """
    Test permissions - members allowed to edit, delete keys but not allowed to create keys
    """
    async with aiohttp.ClientSession() as session:
        master_key = LITELLM_MASTER_KEY

        # Create a team with specific member permissions
        team_data = await create_team(
            session=session,
            key=master_key,
            member_permissions=["/key/update", "/key/delete", "/key/info"],
        )
        team_id = team_data["team_id"]

        # create a user in team=team_id
        user_data = await create_user(
            session=session,
            key=master_key,
            user_id=f"user_{uuid.uuid4().hex[:8]}",
            team_id=team_id,
        )
        user_id = user_data["user_id"]

        # Generate an admin key for the team
        admin_key_data = await generate_key(session, master_key, team_id)
        key_id = admin_key_data["key"]

        # Create a user key
        user_key_data = await generate_key(
            session=session, key=master_key, user_id=user_id
        )
        user_key = user_key_data["key"]

        # Test valid permissions
        # User tries editing a key with team_id
        update_result = await update_key(
            session=session, key=user_key, key_id=key_id, team_id=team_id
        )
        assert (
            "status" not in update_result
        ), "User should be able to update keys for team"

        # User tries deleting a key with team_id
        # Note: Even with /key/delete permission, users can only delete keys they own or if they're team admin
        # The delete endpoint checks ownership/team admin status, not just team member permissions
        delete_result = await delete_key(session=session, key=user_key, key_id=key_id)
        assert (
            "status" in delete_result and delete_result["status"] == 403
        ), "User should not be able to delete keys they don't own (even with /key/delete permission, ownership is required)"

        # Test invalid permissions
        # User tries creating a key with team_id
        create_result = await generate_key(
            session=session, key=user_key, team_id=team_id
        )
        assert (
            "status" in create_result and create_result["status"] != 200
        ), "User should not be able to create keys for team"

        # User tries regenerating a key with team_id
        # Note: Even with /key/regenerate permission, users can only regenerate keys they own or if they're team admin
        regenerate_result = await regenerate_key(
            session=session, key=user_key, key_id=key_id, team_id=team_id
        )
        assert (
            "status" in regenerate_result and regenerate_result["status"] == 401
        ), "User should not be able to regenerate keys they don't own (even with /key/regenerate permission, ownership is required)"


@pytest.mark.asyncio()
async def test_create_permissions():
    """
    Test permissions - members allowed to create keys but not allowed to edit, delete keys
    """
    async with aiohttp.ClientSession() as session:
        master_key = LITELLM_MASTER_KEY

        # Create a team with specific member permissions
        team_data = await create_team(
            session=session, key=master_key, member_permissions=["/key/generate"]
        )
        team_id = team_data["team_id"]

        # Create a user in the team
        user_id = f"user_{uuid.uuid4().hex[:8]}"
        await add_team_member(
            session=session,
            key=master_key,
            team_id=team_id,
            user_id=user_id,
            role="user",
        )

        # Generate an admin key for the team
        admin_key_data = await generate_key(
            session=session, key=master_key, team_id=team_id
        )
        admin_key = admin_key_data["key"]
        key_id = admin_key_data["key"]

        # Create a user key
        user_key_data = await generate_key(
            session=session, key=master_key, user_id=user_id
        )
        user_key = user_key_data["key"]

        # Test valid permissions
        # User tries creating a key with team_id
        create_result = await generate_key(
            session=session, key=user_key, team_id=team_id
        )
        print("success, user created key for team=", create_result)
        assert "key" in create_result, "User should be able to create keys for team"
        assert (
            create_result["team_id"] == team_id
        ), "User should be able to create keys for team"
        assert (
            "status" not in create_result
        ), "User should be able to create keys for team"

        # Test invalid permissions
        # User tries editing a key with team_id
        update_result = await update_key(
            session=session, key=user_key, key_id=key_id, team_id=team_id
        )
        assert (
            "status" in update_result and update_result["status"] != 200
        ), "User should not be able to update keys for team"

        # User tries deleting a key with team_id
        delete_result = await delete_key(session=session, key=user_key, key_id=key_id)
        assert (
            "status" in delete_result and delete_result["status"] == 403
        ), "User should not be able to delete keys for team"

        # User tries regenerating a key with team_id
        # User doesn't have /key/regenerate permission, so should get 401 (team member permission error)
        regenerate_result = await regenerate_key(
            session=session, key=user_key, key_id=key_id, team_id=team_id
        )
        assert (
            "status" in regenerate_result and regenerate_result["status"] == 401
        ), "User should not be able to regenerate keys for team (no /key/regenerate permission)"
        error_data = json.loads(regenerate_result["error"])
        assert (
            error_data["error"]["type"]
            == ProxyErrorTypes.team_member_permission_error.value
        ), "Error should be a team member permission error"