from flask import Flask, request, jsonify, Response, stream_with_context
import requests
import uuid
import json
import time
import os
import re
import base64

API_KEY = os.getenv("API_KEY", "linux.do")

app = Flask(__name__)

@app.before_request
def check_api_key():
    key = request.headers.get("Authorization")
    if key != "Bearer "+API_KEY:
        return jsonify({"success": False, "message": "Unauthorized: Invalid API key"}), 403

@app.route('/v1/models', methods=['GET'])
def get_models():
     # 构建请求头
    headers = {
        "Content-Type": "application/json",
        "Authorization": "Bearer "+API_KEY
    }
    # 发送请求到Akash API
    response = requests.get(
        'https://text.pollinations.ai/models',
        headers=headers
    )
    # 定义名称映射
    name_mapping = {
        "claude-hybridspace": "claude hybridspace",
        "gemini-thinking": "gemini thinking",
        "gemini": "openai/gemini"
    }

    models_data = response.json()
    current_timestamp = int(time.time())   
    converted_data = {
        "object": "list",
        "data": [
            {
                "id": name_mapping.get(model["name"], model["name"]),
                "object": "model",
                "created": current_timestamp,
                "owned_by": "openai" if "openai" in model["name"] else "third_party",
                "permissions": [],
                "root": name_mapping.get(model["name"], model["name"]),
                "parent": None,
                "capabilities": {
                    "vision": model.get("vision", False),
                    "audio": model.get("audio", False),
                    "reasoning": model.get("reasoning", True)
                },
                "description": model["description"]
            }
            for model in models_data
        ]
    }
    return converted_data


def encode_audio_base64(file_path):
    """ Reads an audio file and encodes it as base64. """
    with open(file_path, "rb") as audio_file:
        return base64.b64encode(audio_file.read()).decode("utf-8")

def transcribe_audio(base64_audio, file_format, model="openai-audio"):
    """ Sends base64 encoded audio for transcription using OpenAI-compatible API. """
    API_URL = "https://text.pollinations.ai/openai"  # OpenAI 兼容 API 地址
    API_KEY = "your_api_key_here"
    
    headers = {"Authorization": f"Bearer {API_KEY}", "Content-Type": "application/json"}
    payload = {
        "model": model,
        "messages": [
            {"role": "user", "content": [
                {"type": "text", "text": "Transcribe this audio exactly"},
                {"type": "input_audio", "input_audio": {"data": base64_audio, "format": file_format}}
            ]}
        ]
    }
    
    response = requests.post(API_URL, json=payload, headers=headers)
    print(response.text)
    if response.ok:
        return response.json().get("choices", [{}])[0].get("message", {}).get("content", "No transcription found.")
    return None

@app.route("/v1/audio/transcriptions", methods=["POST"])
def handle_transcription():
    """ Flask API endpoint to handle audio transcription requests. """
    if "file" not in request.files:
        return jsonify({"error": "No file provided"}), 400
    
    file = request.files["file"]
    model = request.form.get("model", "openai-audio")
    print(model)
    file_format = "mp3" if file.filename.lower().endswith(".mp3") else "wav"
    base64_audio = base64.b64encode(file.read()).decode("utf-8")
    
    transcription = transcribe_audio(base64_audio, file_format,model)
    
    if transcription:
        return jsonify({"text": transcription})
    return jsonify({"error": "Transcription failed"}), 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5200, debug=True)