"""
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 time
import zipfile
import traceback
import logging
import httpx
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,
                     MAX_FILE_BYTES, pause_ttl, resume_ttl,
                     request_cancel, get_cancel_event)

api_bp = Blueprint("api", __name__, url_prefix="/api")
log = logging.getLogger("routes.api_bp")


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:
        log.warning("upload: no files, sid=%s", sid)
        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
        data = f.read()
        log.info("upload: sid=%s file=%r size=%d", sid, f.filename, len(data))
        if len(data) > MAX_FILE_BYTES:
            log.warning("upload: file exceeds %dMB, skipped sid=%s file=%r size=%d",
                        MAX_FILE_BYTES // (1024 * 1024), sid, f.filename, len(data))
            continue
        if not add_file(sid, f.filename, data):
            log.warning("upload: session not found/limit, sid=%s file=%r", sid, f.filename)
            return jsonify({"ok": False, "error": "Session not found"}), 404
        added += 1

    if added == 0:
        err = "No filename" if had_unnamed else "No file"
        log.warning("upload: %s, sid=%s", err, sid)
        return jsonify({"ok": False, "error": err}), 400

    log.info("upload: done sid=%s added=%d total=%d", sid, added, file_count(sid))
    return jsonify({"ok": True, "session": sid, "count": file_count(sid)})


@api_bp.route("/upload_refs", methods=["POST"])
def upload_refs():
    """Принять ссылки на файлы (загружены на ВМ-буфер), забрать по egress.

    Вход: JSON {"session": "...", "files": [{"name": str, "size": int, "url": str}]}.
    Каждый файл тянется ИСХОДЯЩИМ GET'ом с ВМ (egress не ограничен шлюзом),
    читается по частям (stream), кладётся в сессию. После успешного pull файл
    удаляется с ВМ (best-effort; TTL-чистка на ВМ тоже есть).
    """
    data = request.get_json(silent=True) or {}
    sid = data.get("session") or create_session()
    refs = data.get("files") or []
    if not refs:
        log.warning("upload_refs: no files, sid=%s", sid)
        return jsonify({"ok": False, "error": "No files"}), 400
    added = 0
    try:
        with httpx.Client(timeout=120, follow_redirects=True) as client:
            for ref in refs:
                name = ref.get("name")
                url = ref.get("url")
                if not name or not url:
                    continue
                # Лимит на один файл (50 МБ): сверх лимита — пропускаем (не участвует)
                if (ref.get("size") or 0) > MAX_FILE_BYTES:
                    log.warning("upload_refs: file exceeds %dMB, skip sid=%s file=%r size=%s",
                                MAX_FILE_BYTES // (1024 * 1024), sid, name, ref.get("size"))
                    try:
                        client.delete(url)
                    except Exception:
                        pass
                    continue
                with client.stream("GET", url) as resp:
                    resp.raise_for_status()
                    content = b"".join(resp.iter_bytes())
                log.info("upload_refs: pulled sid=%s file=%r size=%d", sid, name, len(content))
                if len(content) > MAX_FILE_BYTES:
                    log.warning("upload_refs: pulled file exceeds %dMB, skip sid=%s file=%r size=%d",
                                MAX_FILE_BYTES // (1024 * 1024), sid, name, len(content))
                    try:
                        client.delete(url)
                    except Exception:
                        pass
                    continue
                if not add_file(sid, name, content):
                    log.warning("upload_refs: session not found/limit, sid=%s file=%r", sid, name)
                    return jsonify({"ok": False, "error": "Session not found"}), 404
                try:
                    client.delete(url)  # убрать файл с ВМ после загрузки
                except Exception:
                    pass
                added += 1
    except Exception as e:
        log.error("upload_refs: pull error sid=%s: %r", sid, e)
        return jsonify({"ok": False, "error": "Pull failed: %s" % e}), 502
    log.info("upload_refs: done sid=%s added=%d total=%d", sid, added, file_count(sid))
    return jsonify({"ok": True, "session": sid, "count": file_count(sid)})


@api_bp.route("/session_files/<sid>", methods=["GET"])
def session_files(sid):
    """Отладка: список файлов сессии с размерами (для теста загрузки)."""
    files = get_files(sid)
    if files is None:
        return jsonify({"ok": False, "error": "Session not found"}), 404
    return jsonify({
        "ok": True, "count": len(files),
        "files": [{"name": n, "size": len(b)} for n, b in files],
    })


@api_bp.route("/cancel/<sid>", methods=["POST"])
def cancel(sid):
    """Запросить мягкое прерывание обработки сессии.

    Воркер останавливается на ближайшей границе файла (текущий добирается),
    собирает частичный результат (готовые файлы + mapping) и шлёт SSE-событие `cancelled`.
    """
    if request_cancel(sid):
        log.info("cancel: requested sid=%s", sid)
        return jsonify({"ok": True}), 200
    log.warning("cancel: session not found sid=%s", sid)
    return jsonify({"ok": False, "error": "Session not found"}), 404


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

    Все файлы обрабатываются ЕДИНЫМ вызовом obfuscate_files (общий mapping,
    согласованные токены). Обработка идёт в отдельном потоке; прогресс
    передаётся через очередь. Разрыв соединения клиента корректно
    перехватывается и останавливает генератор (и воркер — через cancel_event).
    """
    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]
    log.info("process_stream: start sid=%s files=%d", sid, len(all_files))
    # Сессия живёт, пока идёт обработка (TTL возобновляется в finally генератора)
    pause_ttl(sid)

    def generate():
        llm = LLMClient()
        q = queue.Queue()
        cancel = threading.Event()          # локальный: разрыв клиента (стоп heartbeat)
        cancel_event = get_cancel_event(sid)  # из сессии: мягкая отмена (кнопка «Прервать»)

        # Состояние для глобальной ETA (символы)
        eta = {"total_chars": 0, "done_chars": 0, "cur_chars": 0, "cur_total": 0, "cur_done": 0}
        _TOKENS_PER_CHAR = 8.0  # эмпирический коэффициент символов -> токенов LLM

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

        # Состояние для per-file ETA (чанки)
        fstate = {"t0": 0.0, "total": 0, "done": 0}

        def file_progress(event, fname, **fields):
            if event == "file_start":
                fstate["t0"] = time.time()
                fstate["total"] = fields.get("chunks", 0)
                fstate["done"] = 0
            elif event == "file_chunk":
                fstate["done"] = fields.get("chunks_done", 0)
                elapsed = time.time() - fstate["t0"]
                rate = fstate["done"] / elapsed if elapsed > 0 else 0
                rem = fstate["total"] - fstate["done"]
                if rate > 0:
                    fields["eta_sec"] = max(0, round(rem / rate))
            q.put(("file", event, fname, fields))

        def worker():
            log.info("worker: start sid=%s files=%d", sid, len(all_files))
            t0 = datetime.utcnow()
            try:
                zip_data, csv_str, meta = obfuscate_files(
                    all_files, llm_client=llm, progress_cb=progress,
                    cancel_event=cancel_event, file_progress=file_progress,
                )
                stats = {
                    "tokens": llm.tokens_total,
                    "llm_sec": round(llm.llm_sec, 1),
                }
                dt = (datetime.utcnow() - t0).total_seconds()
                if meta and meta.get("cancelled"):
                    log.info("worker: cancelled sid=%s in %.1fs processed=%d/%d",
                             sid, dt, meta.get("processed", 0), meta.get("total", 0))
                    q.put(("cancelled", zip_data, csv_str, stats, meta))
                else:
                    log.info("worker: done sid=%s in %.1fs tokens=%d llm_sec=%.1f zip_len=%d",
                             sid, dt, llm.tokens_total, llm.llm_sec, len(zip_data))
                    q.put(("result", zip_data, csv_str, stats))
            except Exception as e:
                log.error("worker: exception sid=%s: %r\n%s", sid, e, traceback.format_exc())
                q.put(("error", repr(e)))

        threading.Thread(target=worker, daemon=True).start()
        log.debug("process_stream: worker thread started sid=%s", sid)

        try:
            while True:
                try:
                    evt = q.get(timeout=1)
                except queue.Empty:
                    if cancel.is_set():
                        log.info("process_stream: cancelled sid=%s (generator exit)", sid)
                        return
                    # Heartbeat: живая статистика LLM + глобальная ETA
                    tokens = llm.tokens_total
                    elapsed = llm.llm_elapsed_now()
                    eta_sec = None
                    if eta["total_chars"] > 0 and tokens > 0 and elapsed > 0:
                        est_total_tokens = eta["total_chars"] * _TOKENS_PER_CHAR
                        rate = tokens / elapsed
                        if rate > 0:
                            eta_sec = max(0, int((est_total_tokens - tokens) / rate))
                    done_chars = eta["done_chars"]
                    if eta["cur_total"] > 0:
                        done_chars += int(eta["cur_chars"] * (eta["cur_done"] / eta["cur_total"]))
                    try:
                        yield (
                            f"event: llm\n"
                            f"data: {json.dumps({'active': llm.llm_active, 'elapsed': round(elapsed, 1), 'tokens': tokens, 'eta_sec': eta_sec, 'done_chars': done_chars, 'total_chars': eta['total_chars']})}\n\n"
                        )
                    except _disconnect_exceptions() as e:
                        log.warning("process_stream: disconnect during heartbeat sid=%s err=%r", sid, e)
                        cancel.set()
                        if cancel_event:
                            cancel_event.set()
                        return
                    continue

                kind = evt[0]

                if kind == "progress":
                    _, phase, idx, name, total_, elapsed = evt
                    log.debug("process_stream: event=%s idx=%d name=%r elapsed=%s sid=%s",
                              phase, idx, name, elapsed, sid)
                    try:
                        yield (
                            f"event: {phase}\n"
                            f"data: {json.dumps({'idx': idx, 'name': name, 'total': total_, 'elapsed': elapsed})}\n\n"
                        )
                    except _disconnect_exceptions() as e:
                        log.warning("process_stream: disconnect on progress sid=%s phase=%s err=%r", sid, phase, e)
                        cancel.set()
                        if cancel_event:
                            cancel_event.set()
                        return

                elif kind == "file":
                    _, event, fname, fields = evt
                    # Обновляем состояние для глобальной ETA
                    if event == "extract_done":
                        eta["total_chars"] = fields.get("total_chars", 0)
                    elif event == "file_start":
                        eta["cur_chars"] = fields.get("chars", 0)
                        eta["cur_total"] = fields.get("chunks", 0)
                        eta["cur_done"] = 0
                    elif event == "file_chunk":
                        eta["cur_done"] = fields.get("chunks_done", 0)
                    elif event == "file_done":
                        eta["done_chars"] += eta["cur_chars"]
                        eta["cur_chars"] = 0
                        eta["cur_total"] = 0
                        eta["cur_done"] = 0
                    try:
                        yield (
                            f"event: {event}\n"
                            f"data: {json.dumps({'name': fname, **fields})}\n\n"
                        )
                    except _disconnect_exceptions() as e:
                        log.warning("process_stream: disconnect on file event sid=%s err=%r", sid, e)
                        cancel.set()
                        if cancel_event:
                            cancel_event.set()
                        return

                elif kind == "result":
                    _, zip_data, csv_str, stats = evt
                    log.info("process_stream: result sid=%s, storing result", sid)
                    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"])
                    log.info("process_stream: complete sid=%s count=%d stats=%r", sid, count, stats)
                    try:
                        yield (
                            f"event: complete\n"
                            f"data: {json.dumps({'total': count, **stats})}\n\n"
                        )
                    except _disconnect_exceptions() as e:
                        log.warning("process_stream: disconnect on complete sid=%s err=%r", sid, e)
                        return
                    return

                elif kind == "cancelled":
                    _, zip_data, csv_str, stats, meta = evt
                    log.info("process_stream: cancelled sid=%s, storing partial result", sid)
                    store_result(sid, zip_data)
                    if csv_str:
                        store_csv(sid, csv_str)
                    try:
                        yield (
                            f"event: cancelled\n"
                            f"data: {json.dumps({'saved': meta.get('processed', 0), 'total': meta.get('total', 0), **stats})}\n\n"
                        )
                    except _disconnect_exceptions() as e:
                        log.warning("process_stream: disconnect on cancelled sid=%s err=%r", sid, e)
                        return
                    return

                elif kind == "error":
                    _, msg = evt
                    log.error("process_stream: error event sid=%s msg=%r", sid, msg)
                    try:
                        yield f"event: error\ndata: {json.dumps({'error': msg})}\n\n"
                    except _disconnect_exceptions() as e:
                        log.warning("process_stream: disconnect on error sid=%s err=%r", sid, e)
                        return
                    return
        finally:
            resume_ttl(sid)
            log.debug("process_stream: generator exit sid=%s, TTL resumed", sid)

    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]
        pause_ttl(sid)
        try:
            zip_data, csv_str, _ = obfuscate_files(all_files, llm_client=llm)
        finally:
            resume_ttl(sid)
        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")