import time
import json
import base64
from flask import Flask, render_template, request
from flask_sock import Sock

app = Flask(__name__)
sock = Sock(app)

VERSION = '1.0.18'


@app.route('/')
def index():
    return render_template('index.html', version=VERSION)


@app.route('/health')
def health():
    return {'ok': True, 'version': VERSION}


@sock.route('/ws-upload')
def ws_upload(ws):
    """WebSocket: получает чанки, считает байты и время"""
    t0 = time.time()
    total = 0
    while True:
        chunk = ws.receive()
        if chunk is None:
            break
        if isinstance(chunk, str) and chunk == 'DONE':
            break
        total += len(chunk)
    elapsed = round((time.time() - t0) * 1000)
    ws.send(json.dumps({'ok': True, 'bytes': total, 'elapsed_ms': elapsed}))


@app.route('/upload', methods=['POST'])
def upload():
    t0 = time.time()
    total = 0
    # base64-режим: клиент отправляет поле "data"
    data = request.form.get('data', '')
    if data:
        # убрать префикс data:...;base64,
        if ',' in data:
            data = data.split(',', 1)[1]
        # form-encode превращает + в пробелы — восстановить
        data = data.replace(' ', '+')
        # добавить padding если потерялся
        missing = len(data) % 4
        if missing:
            data += '=' * (4 - missing)
        raw = base64.b64decode(data)
        total = len(raw)
    # файловый режим
    elif request.files:
        for f in request.files.values():
            while True:
                chunk = f.read(8192)
                if not chunk:
                    break
                total += len(chunk)
    # raw-режим: тело как есть
    else:
        total = len(request.get_data())
    elapsed = round((time.time() - t0) * 1000)
    return {'ok': True, 'bytes': total, 'elapsed_ms': elapsed}


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