"""
OpenAI Audio API Proxy Server to Pollinations
Forwards OpenAI-style audio requests to Pollinations.ai API

Usage:
    export POLLINATIONS_TOKEN="your_token_here"
    export ALLOWED_API_KEYS="key1,key2,key3"  # Optional - comma separated keys
    python mock_openai_audio_api.py

Environment Variables:
    POLLINATIONS_TOKEN - Your Pollinations API token (required for forwarding)
    ALLOWED_API_KEYS - Comma-separated list of allowed API keys for auth (optional)
                       If not set, any Bearer token is accepted
"""

from flask import Flask, request, jsonify, send_file, Response
import requests
import io
import base64
import json
import os
from datetime import datetime

app = Flask(__name__)

# Configuration from environment
POLLINATIONS_TOKEN = os.getenv("POLLINATIONS_TOKEN")
ALLOWED_API_KEYS = os.getenv("ALLOWED_API_KEYS", "").split(",") if os.getenv("ALLOWED_API_KEYS") else None
POLLINATIONS_API_URL = "https://gen.pollinations.ai/v1/chat/completions"

# Supported voices and models
SUPPORTED_VOICES = ["alloy", "ash", "ballad", "coral", "echo", "fable", 
                    "onyx", "nova", "sage", "shimmer", "verse", "marin", "cedar"]
TTS_MODELS = ["tts-1", "tts-1-hd", "gpt-4o-mini-tts", "gpt-4o-mini-tts-2025-12-15"]

def validate_api_key():
    """Validate API key from Authorization header"""
    auth_header = request.headers.get('Authorization', '')
    if not auth_header.startswith('Bearer '):
        return False, "Missing or invalid Authorization header"
    
    api_key = auth_header[7:]  # Remove 'Bearer '
    
    # If ALLOWED_API_KEYS is set, validate against the list
    if ALLOWED_API_KEYS:
        if api_key not in ALLOWED_API_KEYS:
            return False, "Invalid API key"
    else:
        # If no specific keys are set, just check it's not empty
        if not api_key:
            return False, "Invalid API key"
    
    return True, None

@app.route('/v1/audio/speech', methods=['POST'])
def create_speech():
    """Text-to-Speech endpoint - forwards to Pollinations"""
    # Validate API key
    is_valid, error_msg = validate_api_key()
    if not is_valid:
        return jsonify({"error": {"message": error_msg, "type": "invalid_request_error"}}), 401
    
    data = request.get_json()
    
    # Validate required parameters
    if not data:
        return jsonify({"error": {"message": "Request body is required", "type": "invalid_request_error"}}), 400
    
    input_text = data.get('input')
    model = data.get('model')
    voice = data.get('voice')
    
    if not input_text:
        return jsonify({"error": {"message": "Missing required parameter: input", "type": "invalid_request_error"}}), 400
    if not model:
        return jsonify({"error": {"message": "Missing required parameter: model", "type": "invalid_request_error"}}), 400
    if not voice:
        return jsonify({"error": {"message": "Missing required parameter: voice", "type": "invalid_request_error"}}), 400
    
    # Validate input length
    if len(input_text) > 4096:
        return jsonify({"error": {"message": "Input text exceeds maximum length of 4096 characters", "type": "invalid_request_error"}}), 400
    
    # Validate model
    if model not in TTS_MODELS:
        return jsonify({"error": {"message": f"Invalid model: {model}", "type": "invalid_request_error"}}), 400
    
    # Validate voice
    voice_id = voice
    if isinstance(voice, str):
        if voice not in SUPPORTED_VOICES:
            return jsonify({"error": {"message": f"Invalid voice: {voice}", "type": "invalid_request_error"}}), 400
    elif isinstance(voice, dict):
        voice_id = voice.get('id', 'alloy')
    
    # Optional parameters
    response_format = data.get('response_format', 'mp3')
    speed = data.get('speed', 1.0)
    instructions = data.get('instructions', '')
    stream_format = data.get('stream_format', 'audio')
    
    # Build system prompt with emotion/instructions
    emotion = instructions if instructions else "neutral"
    system_instruction = f"Only repeat what I say. Now say with proper emphasis in a \"{emotion}\" emotion this statement."
    
    # Prepare Pollinations API request
    pollinations_headers = {
        "Content-Type": "application/json",
    }
    
    # Add Pollinations token if available
    if POLLINATIONS_TOKEN:
        pollinations_headers["Authorization"] = f"Bearer {POLLINATIONS_TOKEN}"
    
    pollinations_payload = {
        "model": "openai-audio",
        "modalities": ["text", "audio"],
        "audio": {
            "voice": voice_id if isinstance(voice_id, str) else voice_id,
            "format": response_format
        },
        "messages": [
            {"role": "system", "content": system_instruction},
            {"role": "user", "content": input_text}
        ]
    }
    
    try:
        # Forward request to Pollinations
        response = requests.post(
            POLLINATIONS_API_URL,
            headers=pollinations_headers,
            json=pollinations_payload,
            timeout=60
        )
        
        if response.status_code != 200:
            # Handle Pollinations errors
            error_message = f"Pollinations API error: {response.status_code}"
            if response.status_code == 402:
                error_message = "Rate limit exceeded. Please try again later or use a premium API key."
            elif response.status_code == 429:
                error_message = "Too many requests. Please slow down."
            elif response.status_code == 401:
                error_message = "Invalid Pollinations token."
            
            return jsonify({
                "error": {
                    "message": error_message,
                    "type": "api_error",
                    "pollinations_status": response.status_code
                }
            }), response.status_code
        
        # Extract audio from Pollinations response
        pollinations_data = response.json()
        try:
            audio_b64 = pollinations_data['choices'][0]['message']['audio']['data']
            audio_bytes = base64.b64decode(audio_b64)
        except (KeyError, IndexError) as e:
            return jsonify({
                "error": {
                    "message": "Invalid response from Pollinations API",
                    "type": "api_error"
                }
            }), 500
        
        # Return audio in OpenAI format
        return send_file(
            io.BytesIO(audio_bytes),
            mimetype=f'audio/{response_format}',
            as_attachment=True,
            download_name=f'speech.{response_format}'
        )
        
    except requests.exceptions.RequestException as e:
        return jsonify({
            "error": {
                "message": f"Network error: {str(e)}",
                "type": "api_error"
            }
        }), 503

@app.route('/v1/audio/transcriptions', methods=['POST'])
def create_transcription():
    """Speech-to-Text endpoint - returns mock data (not yet implemented for Pollinations)"""
    # Validate API key
    is_valid, error_msg = validate_api_key()
    if not is_valid:
        return jsonify({"error": {"message": error_msg, "type": "invalid_request_error"}}), 401
    
    return jsonify({
        "error": {
            "message": "Transcription endpoint not yet implemented. This proxy currently only supports text-to-speech.",
            "type": "not_implemented_error"
        }
    }), 501

@app.route('/v1/audio/translations', methods=['POST'])
def create_translation():
    """Audio Translation endpoint - returns mock data (not yet implemented)"""
    is_valid, error_msg = validate_api_key()
    if not is_valid:
        return jsonify({"error": {"message": error_msg, "type": "invalid_request_error"}}), 401
    
    return jsonify({
        "error": {
            "message": "Translation endpoint not yet implemented. This proxy currently only supports text-to-speech.",
            "type": "not_implemented_error"
        }
    }), 501

@app.route('/v1/audio/voices', methods=['POST'])
def create_voice():
    """Create custom voice endpoint - not supported"""
    return jsonify({
        "error": {
            "message": "Custom voices not supported by Pollinations proxy",
            "type": "not_implemented_error"
        }
    }), 501

@app.route('/v1/audio/voice_consents', methods=['POST', 'GET'])
def voice_consents():
    """Voice consents endpoint - not supported"""
    return jsonify({
        "error": {
            "message": "Voice consents not supported by Pollinations proxy",
            "type": "not_implemented_error"
        }
    }), 501

@app.route('/', methods=['GET'])
def index():
    """Main page showing API status"""
    base_url = request.host_url.rstrip('/')
    
    html = f"""
    <!DOCTYPE html>
    <html>
    <head>
        <title>OpenAI Audio API Proxy</title>
        <style>
            body {{
                font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
                max-width: 800px;
                margin: 50px auto;
                padding: 20px;
                background: #f5f5f5;
            }}
            .container {{
                background: white;
                padding: 30px;
                border-radius: 10px;
                box-shadow: 0 2px 10px rgba(0,0,0,0.1);
            }}
            h1 {{
                color: #10a37f;
                margin-top: 0;
            }}
            .status {{
                background: #e8f5e9;
                border-left: 4px solid #4caf50;
                padding: 15px;
                margin: 20px 0;
                border-radius: 4px;
            }}
            .base-url {{
                background: #f5f5f5;
                padding: 10px;
                border-radius: 4px;
                font-family: monospace;
                font-size: 14px;
                margin: 10px 0;
            }}
            .endpoint {{
                background: #f9f9f9;
                padding: 10px;
                margin: 5px 0;
                border-radius: 4px;
                border-left: 3px solid #10a37f;
            }}
            .config {{
                margin: 20px 0;
            }}
            .config-item {{
                padding: 8px 0;
                border-bottom: 1px solid #eee;
            }}
            .badge {{
                display: inline-block;
                padding: 4px 8px;
                border-radius: 4px;
                font-size: 12px;
                font-weight: bold;
            }}
            .badge-success {{ background: #4caf50; color: white; }}
            .badge-warning {{ background: #ff9800; color: white; }}
            code {{
                background: #f5f5f5;
                padding: 2px 6px;
                border-radius: 3px;
                font-size: 13px;
            }}
        </style>
    </head>
    <body>
        <div class="container">
            <h1>🎙️ OpenAI Audio API Proxy</h1>
            
            <div class="status">
                <strong>✅ API is running at:</strong>
                <div class="base-url">{base_url}</div>
            </div>
            
            <h2>📋 Configuration</h2>
            <div class="config">
                <div class="config-item">
                    <strong>Pollinations Token:</strong> 
                    {"<span class='badge badge-success'>✓ Configured</span>" if POLLINATIONS_TOKEN else "<span class='badge badge-warning'>⚠ Not Set</span>"}
                </div>
                <div class="config-item">
                    <strong>Authentication:</strong> 
                    {"<span class='badge badge-success'>Restricted</span>" if ALLOWED_API_KEYS else "<span class='badge badge-warning'>Open (any token)</span>"}
                </div>
            </div>
            
            <h2>🚀 Available Endpoints</h2>
            <div class="endpoint">
                <strong>POST</strong> <code>/v1/audio/speech</code> - Text-to-Speech
            </div>
            <div class="endpoint">
                <strong>GET</strong> <code>/health</code> - Health check
            </div>
            
            <h2>📝 Example Usage</h2>
            <pre style="background: #f5f5f5; padding: 15px; border-radius: 4px; overflow-x: auto;">
curl {base_url}/v1/audio/speech \\
  -H "Authorization: Bearer YOUR_API_KEY" \\
  -H "Content-Type: application/json" \\
  -d '{{
    "model": "gpt-4o-mini-tts",
    "input": "Hello world!",
    "voice": "alloy"
  }}' \\
  --output speech.mp3</pre>
            
            <h2>🎵 Supported Voices</h2>
            <p>{", ".join(SUPPORTED_VOICES)}</p>
            
            <p style="margin-top: 30px; color: #666; font-size: 14px;">
                Powered by <a href="https://pollinations.ai" target="_blank">Pollinations.ai</a>
            </p>
        </div>
    </body>
    </html>
    """
    return html

@app.route('/health', methods=['GET'])
def health_check():
    """Health check endpoint"""
    pollinations_status = "not_checked"
    
    # Check if we can reach Pollinations
    try:
        if POLLINATIONS_TOKEN:
            test_response = requests.get(
                "https://gen.pollinations.ai",
                timeout=5
            )
            pollinations_status = "reachable" if test_response.status_code < 500 else "error"
        else:
            pollinations_status = "no_token"
    except:
        pollinations_status = "unreachable"
    
    return jsonify({
        "status": "ok",
        "timestamp": datetime.now().isoformat(),
        "pollinations_token_configured": POLLINATIONS_TOKEN is not None,
        "pollinations_status": pollinations_status,
        "auth_mode": "key_list" if ALLOWED_API_KEYS else "open",
        "supported_endpoints": ["/v1/audio/speech"]
    })

@app.errorhandler(404)
def not_found(error):
    return jsonify({
        "error": {
            "message": "Not found. Available endpoint: POST /v1/audio/speech",
            "type": "invalid_request_error"
        }
    }), 404

@app.errorhandler(500)
def internal_error(error):
    return jsonify({
        "error": {
            "message": "Internal server error",
            "type": "api_error"
        }
    }), 500

if __name__ == '__main__':
    print("=" * 70)
    print("OpenAI Audio API Proxy to Pollinations.ai")
    print("=" * 70)
    
    # Configuration status
    print("\n📋 Configuration:")
    if POLLINATIONS_TOKEN:
        print(f"  ✅ Pollinations Token: Configured ({POLLINATIONS_TOKEN[:10]}...)")
    else:
        print("  ⚠️  Pollinations Token: NOT SET (requests will use free tier)")
        print("     Set with: export POLLINATIONS_TOKEN='your_token'")
    
    if ALLOWED_API_KEYS:
        print(f"  ✅ Auth: Restricted to {len(ALLOWED_API_KEYS)} API key(s)")
    else:
        print("  ⚠️  Auth: Open (any Bearer token accepted)")
        print("     Set with: export ALLOWED_API_KEYS='key1,key2,key3'")
    
    print("\n🚀 Available endpoints:")
    print("  POST /v1/audio/speech - Text-to-Speech (forwards to Pollinations)")
    print("  GET  /health - Health check")
    
    print("\n📝 Example usage:")
    print("""
    curl http://localhost:5000/v1/audio/speech \\
      -H "Authorization: Bearer YOUR_KEY" \\
      -H "Content-Type: application/json" \\
      -d '{
        "model": "gpt-4o-mini-tts",
        "input": "Hello world!",
        "voice": "alloy"
      }' \\
      --output speech.mp3
    """)
    
    print("\n🌐 Starting server on http://localhost:7860")
    print("   Visit http://localhost:7860 in your browser to see API status")
    print("=" * 70)
    
    app.run(host='0.0.0.0', port=7860, debug=True)