from asgiref.wsgi import WsgiToAsgi
from flask import Flask, request, Response
import requests
import json
import os
import sys
import logging
from helper import create_jwt

app = Flask(__name__)

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(message)s',
    handlers=[
        logging.StreamHandler(sys.stdout)
    ]
)
logger = logging.getLogger(__name__)

@app.route('/chat/completions', methods=['POST'])
async def chat():
    """
    Handle chat completion requests.
    """
    logger.info("Received chat completion request")
    
    # 记录请求信息
    logger.info(f"Request method: {request.method}")
    logger.info(f"Request URL: {request.url}")
    logger.info(f"Request headers: {dict(request.headers)}")
    
    # Get the payload from the request
    payload = request.json
    logger.info(f"Request payload: {json.dumps(payload, indent=2)}")

    # Get the model from the payload, defaulting to "claude-3-5-sonnet-20240620"
    model = payload.get('model', 'claude-3-5-sonnet-20240620')
    logger.info(f"Using model: {model}")

    # Extract GitHub username and Zed user ID from Authorization header
    auth_header = request.headers.get('Authorization')
    if not auth_header or not auth_header.startswith('Bearer '):
        logger.error("Invalid Authorization header")
        return Response('Invalid Authorization header', status=401)

    try:
        github_username, zed_user_id = auth_header[7:].split(',')
        logger.info(f"GitHub username: {github_username}, Zed user ID: {zed_user_id}")
    except ValueError:
        logger.error("Invalid Authorization header format")
        return Response('Invalid Authorization header format', status=401)

    # Prepare the request for the LLM API
    url = "https://llm.zed.dev/completion"
    logger.info(f"LLM API URL: {url}")
    
    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": ""
        }
    }
    logger.info(f"LLM API payload: {json.dumps(llm_payload, indent=2)}")

    jwt = create_jwt(github_username, int(zed_user_id))
    logger.info(f"Generated JWT token: {jwt}")

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

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

    async def generate():
        try:
            logger.info("Sending request to LLM API")
            with requests.post(url, headers=headers, json=llm_payload, stream=True, proxies=proxies, allow_redirects=True) as response:
                logger.info(f"LLM API response status: {response.status_code}")
                logger.info(f"LLM API response headers: {dict(response.headers)}")
                
                if response.status_code == 301:
                    new_location = response.headers.get('Location')
                    logger.warning(f"Received 301 redirect. New location: {new_location}")
                    # 如果需要，可以在这里处理重定向
                
                for chunk in response.iter_content(chunk_size=1024):
                    if chunk:
                        logger.debug(f"Received chunk of size: {len(chunk)} bytes")
                        yield chunk
        except Exception as e:
            logger.error(f"Error during API request: {str(e)}")
            yield str(e).encode()

    logger.info("Returning streaming response")
    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
    logger.info("Starting the application")
    uvicorn.run(asgi_app, host="0.0.0.0", port=8000)