from asgiref.wsgi import WsgiToAsgi
from flask import Flask, request, Response
import requests
import json
import random
import os
import jwt
import uuid
from datetime import datetime, timedelta
import string

app = Flask(__name__)

def get_github_username_zed_userid_list():
    """
    Get GitHub username and Zed user ID from the USER environment variable.

    Returns:
        list: A list containing a single tuple with the GitHub username (str) and Zed user ID (str).
    """
    user_env = os.environ.get('USER', 'default_user,123456')
    try:
        username, user_id = user_env.split(',')
        return [(username.strip(), user_id.strip())]
    except ValueError:
        print("Warning: Invalid format in USER environment variable. Using default values.")
        return [("default_user", "123456")]

def create_jwt(github_user_login: str, user_id: int) -> str:
    """
    Create a JSON Web Token (JWT) for a given GitHub user.

    Args:
        github_user_login (str): The GitHub username of the user.
        user_id (int): The user's ID.

    Returns:
        str: A JWT encoded string containing user information and authentication details.

    Note:
        The token has a lifetime of 1 hour and includes the following claims:
        - iat: Issued at time
        - exp: Expiration time
        - jti: Unique token identifier
        - userId: User's ID
        - githubUserLogin: GitHub username
        - isStaff: Boolean indicating staff status (default: False)
        - hasLlmClosedBetaFeatureFlag: Boolean for LLM closed beta feature (default: False)
        - plan: User's plan (default: "Free")
    """
    LLM_TOKEN_LIFETIME = timedelta(hours=1)
    now = datetime.utcnow()

    payload = {
        "iat": int(now.timestamp()),
        "exp": int((now + LLM_TOKEN_LIFETIME).timestamp()),
        "jti": str(uuid.uuid4()),
        "userId": user_id,
        "githubUserLogin": github_user_login,
        "isStaff": False,
        "hasLlmClosedBetaFeatureFlag": False,
        "plan": "Free"
    }

    return jwt.encode(payload, 'llm-secret', algorithm='HS256')

@app.route('/chat/completions', methods=['POST'])
async def chat():
    """
    Handle chat completion requests.

    This function processes incoming POST requests to the '/chat/completions' endpoint.
    It prepares the payload for the LLM API, generates a JWT for authentication,
    and streams the response from the LLM API back to the client.

    Returns:
        Response: A streaming response containing the LLM API's output.

    Note:
        - The function uses environment variables for proxy configuration and user data.
        - The LLM model defaults to "claude-3-5-sonnet-20240620" if not specified.
    """
    # Get the payload from the request
    payload = request.json

    # Get the model from the payload, defaulting to "claude-3-5-sonnet-20240620"
    model = payload.get('model', 'claude-3-5-sonnet-20240620')

    # Prepare the request for the LLM API
    url = "https://llm.zed.dev/completion?"
    
    llm_payload = {
        "provider": "anthropic",
        "model": model,
        "provider_request": {
            "model": model,
            "max_tokens": payload.get('max_tokens', 8192),
            "temperature": payload.get('temperature', 0),
            "top_p": payload.get('top_p', 0.7),
            "messages": payload['messages'],
            "system": ""
        }
    }

    github_username, zed_user_id = get_github_username_zed_userid_list()[0]
    jwt = create_jwt(github_username, zed_user_id)

    headers = {
        'Host': 'llm.zed.dev',
        'accept': '*/*',
        'content-type': 'application/json',
        'authorization': f'Bearer {jwt}',
        'user-agent': 'Zed/0.149.3 (macos; aarch64)'
    }

    # Get proxy from environment variable
    proxy = os.environ.get('HTTP_PROXY', None)
    proxies = {'http': proxy, 'https': proxy} if proxy else None

    async def generate():
        with requests.post(url, headers=headers, json=llm_payload, stream=True, proxies=proxies) as response:
            for chunk in response.iter_content(chunk_size=1024):
                if chunk:
                    yield chunk

    return Response(generate(), content_type='application/octet-stream')

# Convert the Flask app to an ASGI app
asgi_app = WsgiToAsgi(app)

if __name__ == '__main__':
    import uvicorn
    uvicorn.run(asgi_app, host="0.0.0.0", port=8000)