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

app = Flask(__name__)

GMI_CHAT_API_URL = "https://console.gmicloud.ai/chat"
GMI_MODELS_API_URL = "https://console.gmicloud.ai/api/v1/ie/artifact/get_public_artifacts"

def get_proxies():
    proxy = os.environ.get('HTTP_PROXY') or os.environ.get('HTTPS_PROXY')
    if proxy:
        return {"http": proxy, "https": proxy}
    return None

def require_auth(f):
    @wraps(f)
    def decorated(*args, **kwargs):
        auth_header = request.headers.get('Authorization')
        if not auth_header:
            return jsonify({"error": "Authorization header missing"}), 401
        
        try:
            auth_type, token = auth_header.split()
            if auth_type.lower() != 'bearer':
                return jsonify({"error": "Invalid authorization type"}), 401
        except ValueError:
            return jsonify({"error": "Invalid authorization header format"}), 401

        password = os.environ.get('ADAPTER_PASSWORD', '123456')
        if token != password:
            return jsonify({"error": "Invalid token"}), 401
        
        return f(*args, **kwargs)
    return decorated

@app.route('/v1/models', methods=['GET'])
@require_auth
def get_models():
    try:
        response = requests.get(GMI_MODELS_API_URL, proxies=get_proxies())
        response.raise_for_status()
        models_data = response.json()
        
        formatted_models = {
            "object": "list",
            "data": []
        }
        
        for model in list(models_data):
            if "model_price" in model and model["model_price"] != None and "modelName" in model["model_price"]:
                formatted_models["data"].append({
                    "id": model["model_price"]["modelName"],
                    "object": "model",
                    "created": time.time(),
                    "owned_by": "gmi"
                })
        
        return jsonify(formatted_models)
    except requests.exceptions.RequestException as e:
        return jsonify({"error": str(e)}), 500

@app.route('/v1/chat/completions', methods=['POST'])
@require_auth
def chat_completions():
    data = request.get_json()
    headers = {
        'Content-Type': 'application/json',
        'origin': 'https://console.gmicloud.ai'
    }
    
    stream = data.get('stream', False)

    try:
        response = requests.post(GMI_CHAT_API_URL, headers=headers, json=data, stream=stream, proxies=get_proxies())
        response.raise_for_status()

        if stream:
            iterator = response.iter_lines()
            try:
                first_line = next(iterator)
            except StopIteration:
                return Response(status=204)

            # The first line might be empty, skip it if so
            if not first_line:
                try:
                    first_line = next(iterator)
                except StopIteration:
                    return Response(status=204)

            if not first_line.strip().startswith(b'data:'):
                error_message = first_line.decode('utf-8')
                # Try to get more of the error message if it exists
                try:
                    error_message += "\n" + "\n".join(line.decode('utf-8') for line in iterator)
                except Exception:
                    pass
                return jsonify({
                    "error": {
                        "message": error_message,
                        "type": "invalid_request_error",
                        "param": None,
                        "code": None
                    }
                }), 400

            def generate():
                yield first_line + b'\n'
                for line in iterator:
                    yield line + b'\n'
            
            return Response(generate(), content_type=response.headers['content-type'])
        else:
            try:
                return jsonify(response.json()), response.status_code
            except json.JSONDecodeError:
                error_message = response.text
                return jsonify({
                    "error": {
                        "message": error_message,
                        "type": "invalid_request_error",
                        "param": None,
                        "code": None
                    }
                }), 400
            
    except requests.exceptions.RequestException as e:
        return jsonify({"error": str(e)}), 500

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