import csv
from typing import Optional
import httpx
import json
import asyncio

proxy_base_url = "http://0.0.0.0:4000" # 👈 SET TO PROXY URL
master_key = "sk-1234" # 👈 SET TO PROXY MASTER KEY

## GLOBAL HTTP CLIENT ## - faster http calls
class HTTPHandler:
    def __init__(self, concurrent_limit=1000):
        # Create a client with a connection pool
        self.client = httpx.AsyncClient(
            limits=httpx.Limits(
                max_connections=concurrent_limit,
                max_keepalive_connections=concurrent_limit,
            )
        )

    async def close(self):
        # Close the client when you're done with it
        await self.client.aclose()

    async def get(
        self, url: str, params: Optional[dict] = None, headers: Optional[dict] = None
    ):
        response = await self.client.get(url, params=params, headers=headers)
        return response

    async def post(
        self,
        url: str,
        data: Optional[dict] = None,
        params: Optional[dict] = None,
        headers: Optional[dict] = None,
    ):
        try:
            response = await self.client.post(
                url, data=data, params=params, headers=headers
            )
            return response
        except Exception as e:
            raise e

async def import_sheet():
    tasks = []
    http_client = HTTPHandler()
    with open('my-batch-sheet.csv', 'r') as file:
        csv_reader = csv.DictReader(file)
        for row in csv_reader:
            task = create_user(client=http_client, user_id=row['ID'], max_budget=row['Max Budget'], user_name=row['Name'])
            tasks.append(task)
            # print(f"ID: {row['ID']}, Name: {row['Name']}, Max Budget: {row['Max Budget']}")

    keys = await asyncio.gather(*tasks)

    with open('my-batch-sheet_new.csv', 'w', newline='') as new_file:
        fieldnames = ['ID', 'Name', 'Max Budget', 'keys']
        csv_writer = csv.DictWriter(new_file, fieldnames=fieldnames)
        csv_writer.writeheader()

        with open('my-batch-sheet.csv', 'r') as file:
            csv_reader = csv.DictReader(file)
            for i, row in enumerate(csv_reader):
                row['keys'] = keys[i]  # Add the 'keys' value from the corresponding task result
                csv_writer.writerow(row)

    await http_client.close()

asyncio.run(import_sheet())

async def create_key_with_alias(client: HTTPHandler, user_id: str, max_budget: float):
    global proxy_base_url
    if not proxy_base_url.endswith("/"):
        proxy_base_url += "/"
    url = proxy_base_url + "key/generate"

    # call /key/generate
    print("CALLING /KEY/GENERATE")
    response = await client.post(
        url=url,
        headers={"Authorization": f"Bearer {master_key}"},
        data=json.dumps({
            "user_id": user_id,
            "key_alias": f"{user_id}-key",
            "max_budget": max_budget # 👈 KEY CHANGE: SETS MAX BUDGET PER KEY
        })
    )
    print(f"response: {response.text}")
    return response.json()["key"]

async def create_user(client: HTTPHandler, user_id: str, max_budget: float, user_name: str):
    """
    - call /user/new
    - create key for user
    """
    global proxy_base_url
    if not proxy_base_url.endswith("/"):
        proxy_base_url += "/"
    url = proxy_base_url + "user/new"

    # call /user/new
    await client.post(
        url=url,
        headers={"Authorization": f"Bearer {master_key}"},
        data=json.dumps({
            "user_id": user_id,
            "user_alias": user_name,
            "auto_create_key": False,
            # "max_budget": max_budget # 👈 [OPTIONAL] Sets max budget per user (if you want to set a max budget across keys)
        })
    )

    # create key for user
    return await create_key_with_alias(client=client, user_id=user_id, max_budget=max_budget)