import requests
from flask import Flask, request, Response, jsonify

import os  


app = Flask(__name__)

cookie = os.environ.get('COOKIE')

# Authentication headers (this could be made dynamic too)
headers = {
    'accept': 'application/json',
    'content-type': 'application/json',
    'cookie': cookie
}

# Base URL for proxying the requests
base_url = 'https://openui.fly.dev'

@app.route('/')
def hello():
    return 'Hello, World!'

# /v1/models endpoint
@app.route('/v1/models', methods=['GET','POST'])
def get_models():  
    models_data = {  
          "data": [  
            {  
              "id": "gpt-3.5-turbo",  
              "object": "model",  
              "created": 1,  
              "token_coefficient": 0.5,  
              "max_tokens": 16385  
            },  
            {  
              "id": "gpt-4-turbo",  
              "object": "model",  
              "created": 1,  
              "vision": True,  
              "token_coefficient": 1,  
              "max_tokens": 128000  
            },  
            {  
              "id": "gpt-4o",  
              "object": "model",  
              "created": 1,  
              "vision": True,  
              "token_coefficient": 0.9,  
              "max_tokens": 128000  
            },  
            {  
              "id": "gpt-4o-mini",  
              "object": "model",  
              "created": 1,  
              "vision": True,  
              "token_coefficient": 0.2,  
              "max_tokens": 128000  
            }  
          ]  
        } 

    return jsonify(models_data)

# /v1/chat/completions endpoint with streaming
@app.route('/v1/chat/completions', methods=['POST','GET'])
def chat_completions():
    # Extract request data (JSON payload)
    payload = request.json
    
    # Forward the POST request to the target API with streaming responses
    url = f'{base_url}/v1/chat/completions'
    response = requests.post(url, headers=headers, json=payload, stream=True)
    
    # Streaming response back to the client
    def stream_response():
        for chunk in response.iter_content(chunk_size=1024):
            if chunk:
                yield chunk.decode('utf-8')
    
    return Response(stream_response(), content_type='application/json')

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0')  # To run it in docker container