from flask import Flask, request, Response
from flask_cors import CORS
import requests
import json
from datetime import datetime, timedelta
import os

app = Flask(__name__)
CORS(app)

CHATANYWHERE_ENDPOINT = 'https://api.chatanywhere.tech'

class APIKeyManager:
    def __init__(self):
        self.api_keys = os.environ.get('API_KEYS', '').split(',')
        self.daily_uses = {key: 0 for key in self.api_keys}
        self.last_reset = datetime.now().date()

    def reset_daily_uses(self):
        today = datetime.now().date()
        if today > self.last_reset:
            self.daily_uses = {key: 0 for key in self.api_keys}
            self.last_reset = today

    def get_available_key(self):
        self.reset_daily_uses()
        for key in self.api_keys:
            if self.daily_uses[key] < 200:  # 假设每个密钥每天可用200次
                self.daily_uses[key] += 1
                return key
        return None

    def get_key_info(self):
        self.reset_daily_uses()
        return [{"key": key, "uses": uses} for key, uses in self.daily_uses.items()]

key_manager = APIKeyManager()

@app.route('/hf/v1/models', methods=['GET'])
def list_models():
    key = key_manager.get_available_key() 
    if not key:
        return Response(json.dumps({"error": "No available API keys"}), status=429, mimetype='application/json')

    headers = {
        "Authorization": f"Bearer {key}"
    }

    try:
        response = requests.get(f"{CHATANYWHERE_ENDPOINT}/v1/models", headers=headers)
        if response.status_code == 200:
            return Response(response.text, mimetype='application/json')
        else:
            error_message = f"Error: {response.status_code}, {response.text}"
            return Response(json.dumps({"error": error_message}), status=response.status_code, mimetype='application/json')
    except requests.RequestException as e:
        return Response(json.dumps({"error": str(e)}), status=500, mimetype='application/json')

@app.route('/hf/v1/chat/completions', methods=['POST'])
def proxy_request():
    data = request.get_json()
    messages = data.get('messages', [])
    model = data.get('model', 'gpt-4o-mini')
    max_tokens = data.get('max_tokens', 4096)
    temperature = data.get('temperature', 1)
    stream = data.get('stream', False)
    key = key_manager.get_available_key()
    
    if not key:
        return Response(json.dumps({"error": "No available API keys"}), status=429, mimetype='application/json')

    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {key}"
    }

    payload = {
        "model": model,
        "messages": messages,
        "max_tokens": max_tokens,
        "temperature": temperature,
        "stream": stream
    }

    if stream:
        def generate():
            response = requests.post(f"{CHATANYWHERE_ENDPOINT}/v1/chat/completions", headers=headers, json=payload, stream=True)
            for line in response.iter_lines():
                if line:
                    try:
                        line = line.decode('utf-8')
                        if line.startswith("data: "):
                            json_str = line[6:]  # 去掉 "data: " 前缀
                            if json_str.strip() == "[DONE]":
                                yield f"data: [DONE]\n\n"
                            else:
                                data = json.loads(json_str)
                                if 'choices' in data and len(data['choices']) > 0:
                                    choice = data['choices'][0]
                                    if 'delta' in choice:
                                        delta = choice['delta']
                                        if 'role' in delta:
                                            yield f"data: {json.dumps({'choices': [{'delta': {'role': 'assistant'}}]})}\n\n"
                                        elif 'content' in delta:
                                            yield f"data: {json.dumps({'choices': [{'delta': {'content': delta['content']}}]})}\n\n"
                    except json.JSONDecodeError:
                        continue
        
        return Response(generate(), content_type='text/event-stream')
    else:
        response = requests.post(f"{CHATANYWHERE_ENDPOINT}/v1/chat/completions", headers=headers, json=payload)
        if response.status_code == 200:
            return Response(response.content, mimetype='application/json')
        else:
            return Response(json.dumps({"error": f"Error: {response.status_code} - {response.text}"}), 
                            status=response.status_code, mimetype='application/json')

@app.route('/key_info', methods=['GET'])
def key_info():
    return Response(json.dumps(key_manager.get_key_info()), mimetype='application/json')

if __name__ == '__main__':
    print(os.environ.get('API_KEYS', '').split(','))
    port = int(os.environ.get('PORT', 7860))
    app.run(host='0.0.0.0', port=port)