import os
from flask import Flask, request, jsonify, Response, stream_with_context
from functools import wraps
import requests
import json

app = Flask(__name__)

AZURE_API_KEY = os.environ.get("AZURE_API_KEY")
AZURE_API_BASE = "https://openai-skyline-jp.openai.azure.com"
AZURE_API_VERSION = "2023-05-15"
AZURE_DEPLOYMENT_NAME = "GPT-4"
API_TOKEN = os.environ.get("API_TOKEN")

def token_required(f):
    @wraps(f)
    def decorated(*args, **kwargs):
        token = None
        if 'Authorization' in request.headers:
            auth_header = request.headers['Authorization']
            try:
                token = auth_header.split(" ")[1]
            except IndexError:
                return jsonify({"error": "Invalid token format"}), 401
        if not token:
            return jsonify({"error": "Token is missing"}), 401
        if token != API_TOKEN:
            return jsonify({"error": "Invalid token"}), 401
        return f(*args, **kwargs)
    return decorated

def convert_to_openai_format(azure_response):
    openai_response = {
        "id": azure_response.get("id", ""),
        "object": "chat.completion",
        "created": azure_response.get("created", 0),
        "model": azure_response.get("model", "gpt-4"),
        "choices": [],
        "usage": azure_response.get("usage", {})
    }
    
    for choice in azure_response.get("choices", []):
        openai_choice = {
            "index": choice.get("index", 0),
            "message": choice.get("message", {}),
            "finish_reason": choice.get("finish_reason")
        }
        openai_response["choices"].append(openai_choice)
    
    return {k: v for k, v in openai_response.items() if v is not None}

@app.route('/hf/v1/chat/completions', methods=['POST'])
@token_required
def chat_completions():
    request_data = request.json
    stream = request_data.get('stream', False)

    azure_url = f"{AZURE_API_BASE}/openai/deployments/{AZURE_DEPLOYMENT_NAME}/chat/completions?api-version={AZURE_API_VERSION}"

    headers = {
        "Content-Type": "application/json",
        "api-key": AZURE_API_KEY
    }

    if not stream:
        response = requests.post(azure_url, json=request_data, headers=headers)

        if response.status_code == 200:
            azure_response = response.json()
            openai_response = convert_to_openai_format(azure_response)
            return jsonify(openai_response), 200
        else:
            return jsonify({"error": "Azure OpenAI API request failed", "details": response.text}), response.status_code
    else:
        def generate():
            with requests.post(azure_url, json=request_data, headers=headers, stream=True) as response:
                if response.status_code != 200:
                    yield f"data: {json.dumps({'error': 'Azure OpenAI API request failed', 'details': response.text})}\n\n"
                    return

                for line in response.iter_lines():
                    if line:
                        azure_data = json.loads(line.decode('utf-8').replace('data: ', ''))
                        
                        openai_data = {
                            "id": azure_data.get("id", ""),
                            "object": azure_data.get("object", ""),
                            "created": azure_data.get("created", 0),
                            "model": azure_data.get("model", "gpt-4"),
                            "choices": [
                                {
                                    "index": choice.get("index", 0),
                                    "delta": choice.get("delta", {}),
                                    "finish_reason": choice.get("finish_reason")
                                } for choice in azure_data.get("choices", [])
                            ]
                        }
                        
                        openai_data = {k: v for k, v in openai_data.items() if v is not None}
                        
                        yield f"data: {json.dumps(openai_data)}\n\n"

                yield "data: [DONE]\n\n"

        return Response(stream_with_context(generate()), content_type='text/event-stream')

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860)