"""routes/llm.py — LLM-извлечение + сравнение: один SSE-поток."""

import json as _json
import time as _time

from flask import Response

import db
import extractor
import differ


def process(cid):
    """
    GET /llm/process/<cid> — SSE-поток полного цикла:

      1. Для каждого файла договора: parsed_text → LLM → spec_rows
      2. Сравнение допников через differ.diff()
      3. Результат → SSE-сообщения в браузер
    """
    def generate():
        start_time = _time.time()

        # ── Получить все файлы договора ─────────────────────
        supps, _ = db.query(
            "SELECT s.id as supp_id, s.type, d.id as doc_id, d.filename, d.parsed_text "
            "FROM supplements s JOIN documents d ON s.document_id = d.id "
            "WHERE s.contract_id = %s AND d.parsed_text IS NOT NULL "
            "ORDER BY s.created_at",
            (cid,),
        )
        if not supps:
            yield f"data: {_json.dumps({'type': 'error', 'message': 'Нет распарсенных файлов. Сначала нажмите Парсинг.'})}\n\n"
            return

        supp_rows = [dict(zip(supps["columns"], r)) for r in supps["rows"]]
        extracted = {}  # supp_id → [rows]

        # ── Шаг 1: LLM-извлечение для каждого файла ─────────
        for s in supp_rows:
            # Пропустить уже извлечённые
            existing, _ = db.query(
                "SELECT 1 FROM spec_rows WHERE supplement_id = %s LIMIT 1",
                (s["supp_id"],),
            )
            if existing and existing["rows"]:
                # Загрузить существующие
                rows_res, _ = db.query(
                    "SELECT row_num, name, price, qty, sum, date_start "
                    "FROM spec_rows WHERE supplement_id = %s ORDER BY row_num",
                    (s["supp_id"],),
                )
                rows = [dict(zip(rows_res["columns"], r)) for r in rows_res["rows"]] if rows_res else []
                extracted[s["supp_id"]] = rows
                yield f"data: {_json.dumps({'type': 'extract_skip', 'name': s['filename'], 'reason': 'уже извлечено', 'count': len(rows)})}\n\n"
                continue

            yield f"data: {_json.dumps({'type': 'extract_start', 'name': s['filename']})}\n\n"

            t0 = _time.time()
            result = extractor.extract(s["parsed_text"])

            if "error" in result:
                yield f"data: {_json.dumps({'type': 'extract_error', 'name': s['filename'], 'error': result['error']})}\n\n"
                continue

            rows = result.get("rows", [])
            unresolved = result.get("unresolved", [])
            elapsed = round(_time.time() - t0, 2)

            # Сохранить в БД
            db.execute("DELETE FROM spec_rows WHERE supplement_id = %s", (s["supp_id"],))
            for row in rows:
                db.execute(
                    "INSERT INTO spec_rows (supplement_id, row_num, name, price, qty, sum, date_start) "
                    "VALUES (%s, %s, %s, %s, %s, %s, %s)",
                    (
                        s["supp_id"],
                        row.get("row_num"),
                        row.get("name"),
                        row.get("price"),
                        row.get("qty"),
                        row.get("sum"),
                        row.get("date_start"),
                    ),
                )

            extracted[s["supp_id"]] = rows
            yield f"data: {_json.dumps({'type': 'extract_done', 'name': s['filename'], 'time_s': elapsed, 'count': len(rows), 'unresolved': unresolved})}\n\n"

        # ── Шаг 2: Сравнение ────────────────────────────────
        yield f"data: {_json.dumps({'type': 'diff_start', 'files': len(supp_rows)})}\n\n"

        # Группируем: ищем amendments, если нет — сравниваем первый с остальными
        initial_rows = None
        amendments = []

        for s in supp_rows:
            if s["type"] != "initial":
                amendments.append({
                    "supp_id": s["supp_id"],
                    "filename": s["filename"],
                    "type": s["type"],
                    "rows": extracted.get(s["supp_id"], []),
                })
            elif initial_rows is None:
                # Первый initial — базовый
                initial_rows = extracted.get(s["supp_id"], [])
            else:
                # Последующие initial — как допники (fallback)
                amendments.append({
                    "supp_id": s["supp_id"],
                    "filename": s["filename"],
                    "type": "initial",
                    "rows": extracted.get(s["supp_id"], []),
                })

        all_changes = []

        if initial_rows is not None:
            for am in amendments:
                if not am["rows"]:
                    continue
                d = differ.diff(initial_rows, am["rows"])
                changes = d.get("changes", [])
                summary = d.get("summary", {})
                all_changes.append({
                    "filename": am["filename"],
                    "type": am["type"],
                    "changes": changes,
                    "summary": summary,
                })
                yield f"data: {_json.dumps({'type': 'diff_file', 'name': am['filename'], 'changes': len(changes), 'summary': summary})}\n\n"

        # ── Итог ────────────────────────────────────────────
        total_time = round(_time.time() - start_time, 2)
        yield f"data: {_json.dumps({'type': 'done', 'total_time_s': total_time, 'all_changes': all_changes})}\n\n"

    return Response(generate(), mimetype="text/event-stream")