drhider/site/routes/api_bp.py

"""
Blueprint: API DrHider.

Five endpoints:
- POST /api/upload         — upload one file -> {session_id}
- GET  /api/process_stream/<sid> — SSE: process files, per-file progress
- POST /api/process/<sid>  — process all session files -> {status:"done"} (legacy)
- GET  /api/download/<sid> — download ZIP (with timestamp name)
- GET  /api/csv/<sid>      — download CSV separately
"""

import io
import json
import queue
import threading
import zipfile
import traceback
from datetime import datetime, timedelta
from flask import Blueprint, request, send_file, jsonify, Response, stream_with_context

from drhider import obfuscate_files, LLMClient
from session import (create_session, add_file, get_files, store_result,
                     get_result, store_csv, get_csv, cleanup, file_count)

api_bp = Blueprint("api", __name__, url_prefix="/api")


def _disconnect_exceptions():
    """Исключения, означающие отключение клиента SSE."""
    return (GeneratorExit, BrokenPipeError, ConnectionResetError)


@api_bp.route("/upload", methods=["POST"])
def upload():
    """Upload files to session (один или несколько)."""
    sid = request.form.get("session", "")
    if not sid:
        sid = create_session()
    uploaded = request.files.getlist("files")
    if not uploaded:
        return jsonify({"ok": False, "error": "No file"}), 400

    added = 0
    had_unnamed = False
    for f in uploaded:
        if not f.filename:
            had_unnamed = True
            continue
        if not add_file(sid, f.filename, f.read()):
            return jsonify({"ok": False, "error": "Session not found"}), 404
        added += 1

    if added == 0:
        err = "No filename" if had_unnamed else "No file"
        return jsonify({"ok": False, "error": err}), 400

    return jsonify({"ok": True, "session": sid, "count": file_count(sid)})


@api_bp.route("/process_stream/<sid>", methods=["GET"])
def process_stream(sid):
    """SSE: process all session files, streaming per-file progress.

    Все файлы обрабатываются ЕДИНЫМ вызовом obfuscate_files (общий mapping,
    согласованные токены). Обработка идёт в отдельном потоке; прогресс
    передаётся через очередь. Разрыв соединения клиента корректно
    перехватывается и останавливает генератор.
    """
    files = get_files(sid)
    if files is None:
        return jsonify({"ok": False, "error": "Session not found"}), 404
    if not files:
        return jsonify({"ok": False, "error": "No files"}), 400

    all_files = [(fname, content, "") for fname, content in files]

    def generate():
        llm = LLMClient()
        q = queue.Queue()
        cancel = threading.Event()

        def progress(phase, idx, total_, name, elapsed):
            q.put(("progress", phase, idx, name, total_, elapsed))

        def worker():
            try:
                zip_data, csv_str = obfuscate_files(
                    all_files, llm_client=llm, progress_cb=progress
                )
                stats = {
                    "tokens": llm.tokens_total,
                    "llm_sec": round(llm.llm_sec, 1),
                }
                q.put(("result", zip_data, csv_str, stats))
            except Exception as e:
                q.put(("error", repr(e)))

        threading.Thread(target=worker, daemon=True).start()

        while True:
            try:
                evt = q.get(timeout=1)
            except queue.Empty:
                if cancel.is_set():
                    return
                # Heartbeat: живая статистика LLM (для таймера в UI)
                try:
                    yield (
                        f"event: llm\n"
                        f"data: {json.dumps({'active': llm.llm_active, 'elapsed': round(llm.llm_elapsed_now(), 1), 'tokens': llm.tokens_total})}\n\n"
                    )
                except _disconnect_exceptions():
                    cancel.set()
                    return
                continue

            kind = evt[0]

            if kind == "progress":
                _, phase, idx, name, total_, elapsed = evt
                try:
                    yield (
                        f"event: {phase}\n"
                        f"data: {json.dumps({'idx': idx, 'name': name, 'total': total_, 'elapsed': elapsed})}\n\n"
                    )
                except _disconnect_exceptions():
                    cancel.set()
                    return

            elif kind == "result":
                _, zip_data, csv_str, stats = evt
                store_result(sid, zip_data)
                if csv_str:
                    store_csv(sid, csv_str)
                count = 0
                with zipfile.ZipFile(io.BytesIO(zip_data)) as zf:
                    count = len([n for n in zf.namelist() if n != "mapping.csv"])
                try:
                    yield (
                        f"event: complete\n"
                        f"data: {json.dumps({'total': count, **stats})}\n\n"
                    )
                except _disconnect_exceptions():
                    return
                return

            elif kind == "error":
                _, msg = evt
                try:
                    yield f"event: error\ndata: {json.dumps({'error': msg})}\n\n"
                except _disconnect_exceptions():
                    return
                return

    return Response(
        stream_with_context(generate()),
        content_type="text/event-stream",
        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}
    )


@api_bp.route("/process/<sid>", methods=["POST"])
def process(sid):
    """Process all session files -> ZIP (единым вызовом, общий mapping)."""
    files = get_files(sid)
    if files is None:
        return jsonify({"ok": False, "error": "Session not found"}), 404
    if not files:
        return jsonify({"ok": False, "error": "No files"}), 400
    try:
        llm = LLMClient()
        all_files = [(fname, content, "") for fname, content in files]
        zip_data, csv_str = obfuscate_files(all_files, llm_client=llm)
        store_result(sid, zip_data)
        if csv_str:
            store_csv(sid, csv_str)
        return jsonify({"ok": True, "status": "done"})
    except Exception as e:
        traceback.print_exc()
        return jsonify({"ok": False, "error": str(e)}), 500


@api_bp.route("/download/<sid>", methods=["GET"])
def download(sid):
    """Download result and cleanup session."""
    zip_data = get_result(sid)
    if zip_data is None:
        return jsonify({"ok": False, "error": "Not found"}), 404
    ts = (datetime.now() + timedelta(hours=3)).strftime("%Y-%m-%d_%H-%M-%S")
    return send_file(io.BytesIO(zip_data), mimetype="application/zip",
                     as_attachment=True, download_name=f"drhider_{ts}.zip")


@api_bp.route("/csv/<sid>", methods=["GET"])
def csv_download(sid):
    """Download CSV separately."""
    csv_str = get_csv(sid)
    if csv_str is None:
        return jsonify({"ok": False, "error": "Not found"}), 404
    ts = (datetime.now() + timedelta(hours=3)).strftime("%Y-%m-%d_%H-%M-%S")
    buf = io.BytesIO()
    buf.write('\ufeff'.encode('utf-8') + csv_str.encode('utf-8'))
    buf.seek(0)
    return send_file(buf, mimetype="text/csv",
                     as_attachment=True, download_name=f"mapping_{ts}.csv")