import os
import json
import uuid
from flask import Flask, request, Response, jsonify
import socketio
from aiohttp import web
from aiohttp_socks import ProxyConnector

app = Flask(__name__)
sio = socketio.Client()

# 从环境变量中获取API密钥
API_KEY = os.environ.get('PPLX_KEY')

# 代理设置
proxy_url = os.environ.get('PROXY_URL')
connector = ProxyConnector.from_url(proxy_url) if proxy_url else None

sio_opts = {
    'auth': {
        'jwt': 'anonymous-ask-user',
    },
    'reconnection': False,
    'transports': ['websocket'],
    'extraHeaders': {
        'Cookie': os.environ.get('PPLX_COOKIE'),
        'User-Agent': os.environ.get('USER_AGENT'),
        'Accept': '*/*',
        'priority': 'u=1, i',
        'Referer': 'https://www.perplexity.ai/',
    }
}

def log_request(req, status):
    timestamp = datetime.now().isoformat()
    ip = req.remote_addr
    route = req.path
    print(f"{timestamp} - {ip} - {route} - {status}")

def validate_api_key():
    api_key = request.headers.get('x-api-key')
    if api_key != API_KEY:
        log_request(request, 401)
        return jsonify({"error": "Invalid API key"}), 401
    return None

@app.route('/')
def root():
    log_request(request, 200)
    return jsonify({
        "message": "Welcome to the Perplexity AI Proxy API",
        "endpoints": {
            "/ai/v1/messages": {
                "method": "POST",
                "description": "Send a message to the AI",
                "headers": {
                    "x-api-key": "Your API key (required)",
                    "Content-Type": "application/json"
                },
                "body": {
                    "messages": "Array of message objects",
                    "stream": "Boolean (true for streaming response)",
                }
            }
        }
    })

@app.route('/ai/v1/messages', methods=['POST'])
def messages():
    auth_error = validate_api_key()
    if auth_error:
        return auth_error

    try:
        json_body = request.json
        if not json_body.get('stream', False):
            log_request(request, 200)
            return jsonify({
                "id": str(uuid.uuid4()),
                "content": [
                    {"text": "Please turn on streaming."},
                    {"id": "string", "name": "string", "input": {}}
                ],
                "model": "string",
                "stop_reason": "end_turn",
                "stop_sequence": "string",
                "usage": {"input_tokens": 0, "output_tokens": 0}
            })
        
        def generate():
            previous_messages = "\n\n".join([msg['content'] for msg in json_body['messages']])
            msg_id = str(uuid.uuid4())

            yield create_event("message_start", {
                "type": "message_start",
                "message": {
                    "id": msg_id,
                    "type": "message",
                    "role": "assistant",
                    "content": [],
                    "model": "claude-3-opus-20240229",
                    "stop_reason": None,
                    "stop_sequence": None,
                    "usage": {"input_tokens": 8, "output_tokens": 1},
                },
            })
            yield create_event("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}})
            yield create_event("ping", {"type": "ping"})

            @sio.on('connect')
            def on_connect():
                sio.emit('perplexity_ask', previous_messages, {
                    "version": "2.9",
                    "source": "default",
                    "attachments": [],
                    "language": "en-GB",
                    "timezone": "Europe/London",
                    "search_focus": "writing",
                    "frontend_uuid": str(uuid.uuid4()),
                    "mode": "concise",
                    "is_related_query": False,
                    "is_default_related_query": False,
                    "visitor_id": str(uuid.uuid4()),
                    "frontend_context_uuid": str(uuid.uuid4()),
                    "prompt_source": "user",
                    "query_source": "home"
                })

            @sio.on('query_progress')
            def on_query_progress(data):
                if 'text' in data:
                    text = json.loads(data['text'])
                    chunk = text['chunks'][-1] if text['chunks'] else None
                    if chunk:
                        yield create_event("content_block_delta", {
                            "type": "content_block_delta",
                            "index": 0,
                            "delta": {"type": "text_delta", "text": chunk},
                        })

            @sio.on('disconnect')
            def on_disconnect():
                print(" > [Disconnected]")

            @sio.on('error')
            def on_error(error):
                print(f"Socket error: {error}")
                yield create_event("content_block_delta", {
                    "type": "content_block_delta",
                    "index": 0,
                    "delta": {"type": "text_delta", "text": "Error occurred while fetching output 输出时出现错误\nPlease refer to the log for more information 请查看日志以获取更多信息"},
                })
                yield create_event("content_block_stop", {"type": "content_block_stop", "index": 0})
                yield create_event("message_delta", {
                    "type": "message_delta",
                    "delta": {"stop_reason": "end_turn", "stop_sequence": None},
                    "usage": {"output_tokens": 12},
                })
                yield create_event("message_stop", {"type": "message_stop"})
                log_request(request, 500)

            sio.connect('wss://www.perplexity.ai/', **sio_opts)
            sio.wait()
            sio.disconnect()

            yield create_event("content_block_stop", {"type": "content_block_stop", "index": 0})
            yield create_event("message_delta", {
                "type": "message_delta",
                "delta": {"stop_reason": "end_turn", "stop_sequence": None},
                "usage": {"output_tokens": 12},
            })
            yield create_event("message_stop", {"type": "message_stop"})
            log_request(request, 200)

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

    except Exception as e:
        print(f"Request error: {str(e)}")
        log_request(request, 400)
        return jsonify({"error": str(e)}), 400

@app.errorhandler(404)
def not_found(error):
    log_request(request, 404)
    return "Not Found", 404

@app.errorhandler(500)
def server_error(error):
    print(f"Server error: {str(error)}")
    log_request(request, 500)
    return "Something broke!", 500

def create_event(event, data):
    if isinstance(data, dict):
        data = json.dumps(data)
    return f"event: {event}\ndata: {data}\n\n"

if __name__ == '__main__':
    port = int(os.environ.get('PORT', 8081))
    print(f"Perplexity proxy listening on port {port}")
    if not API_KEY:
        print("Warning: PPLX_KEY environment variable is not set. API key validation will fail.")
    app.run(host='0.0.0.0', port=port)