"""
transcription_service.py — Proses transkripsi + enkripsi hasil + streaming segmen.

Perubahan dari versi sebelumnya:
- FIX BUG: _update_progress() tidak lagi overwrite status — hanya update kolom progress.
  Bug lama: setelah worker selesai set status="completed", _update_progress() yang
  dipanggil terakhir kali bisa menimpa kembali menjadi status="processing".
- Tambah streaming segmen: partial_result dan segments_json diupdate setiap N segmen
  sehingga user bisa melihat hasil sementara tanpa menunggu transkripsi selesai.
- Tambah cancel support: worker cek is_cancelled flag setelah setiap segmen.
"""

import json
import logging
import os

from app.core.database import get_db_context
from app.models.transkripsi import TranskripsiTask
from app.services.whisper_local import transcribe_local, CancelledError
from app.services.whisper_api import transcribe_api
from app.services.task_service import update_task_status, is_task_cancelled
from app.core.config import ENGINE_LOCAL
from app.utils.crypto_utils import encrypt_text
from app.utils.court import apply_court_mode

logger = logging.getLogger(__name__)

# Update partial result ke DB setiap N segmen baru
_SEGMENT_FLUSH_EVERY = 3


def process_task(task_id: int) -> None:
    """
    Proses transkripsi secara synchronous di thread pool.

    Alur status:
        pending → processing (20) → loading model (40) → transcribing (segmen per segmen)
        → completed (100)  /  failed (0)  /  cancelled (0)

    Perubahan penting:
    - Status hanya diubah oleh fungsi ini, bukan _update_progress().
    - Hasil segmen di-flush ke DB setiap _SEGMENT_FLUSH_EVERY segmen.
    - Flag is_cancelled dicek setiap segmen selesai.
    """

    audio_file    = None
    engine        = None
    model_type    = None
    is_court_mode = False

    # ── Ambil task & set status "processing" ───────────────────────
    with get_db_context() as db:
        task: TranskripsiTask | None = db.query(TranskripsiTask).filter(
            TranskripsiTask.id == task_id
        ).first()

        if not task:
            logger.warning("[Worker] Task ID %d tidak ditemukan di DB.", task_id)
            return

        # Cek cancel sebelum mulai
        if task.is_cancelled:
            _set_cancelled(task_id)
            return

        if not task.audio_file or not os.path.exists(task.audio_file):
            logger.error(
                "[Worker] Task %d: file audio tidak ditemukan: %s",
                task_id, task.audio_file,
            )
            update_task_status(
                db, task,
                status        = "failed",
                progress      = 0,
                error_message = "File audio tidak ditemukan di server.",
            )
            return

        audio_file    = task.audio_file
        engine        = task.engine
        model_type    = task.model_type
        is_court_mode = task.is_court_mode or False

        update_task_status(db, task, status="processing", progress=20)
        logger.info(
            "[Worker] Task %d mulai diproses (engine=%s, model=%s)",
            task_id, engine, model_type,
        )

    # ── Update progress: loading model ─────────────────────────────
    _only_update_progress(task_id, 40)

    # ── Transkripsi ────────────────────────────────────────────────
    output        = None
    error_message = None
    cancelled     = False

    try:
        if engine == ENGINE_LOCAL:
            # Siapkan callback untuk streaming segmen
            seg_buffer     = []
            seg_flush_count = [0]  # pakai list agar bisa dimodifikasi di closure

            def on_segments(segs: list[dict]) -> None:
                """
                Dipanggil whisper_local setiap segmen selesai.
                Kita flush ke DB setiap _SEGMENT_FLUSH_EVERY segmen baru.
                """
                new_count = len(segs)
                seg_buffer.clear()
                seg_buffer.extend(segs)

                seg_flush_count[0] += 1
                if seg_flush_count[0] % _SEGMENT_FLUSH_EVERY == 0:
                    _flush_segments(task_id, segs, is_court_mode)

            def check_cancel() -> bool:
                """Buka sesi DB baru untuk baca nilai is_cancelled terkini."""
                with get_db_context() as db:
                    return is_task_cancelled(db, task_id)

            output = transcribe_local(
                audio_file,
                model_type,
                segment_callback = on_segments,
                cancel_check     = check_cancel,
            )

            # Flush segmen terakhir yang mungkin belum di-flush
            if seg_buffer:
                _flush_segments(task_id, seg_buffer, is_court_mode)

        else:
            # Engine API: tidak ada streaming segmen, langsung dapat hasil
            output = transcribe_api(audio_file)

        _only_update_progress(task_id, 80)
        logger.info(
            "[Worker] Task %d selesai ditranskripsi (%.2fs).",
            task_id,
            output.get("processing_time") or 0,
        )

    except CancelledError:
        cancelled = True
        logger.info("[Worker] Task %d dibatalkan saat transkripsi.", task_id)

    except Exception as exc:
        error_message = str(exc)
        logger.error("[Worker] Task %d gagal: %s", task_id, exc, exc_info=True)

    # ── Simpan hasil dalam sesi DB baru ────────────────────────────
    with get_db_context() as db:
        task = db.query(TranskripsiTask).filter(
            TranskripsiTask.id == task_id
        ).first()

        if not task:
            logger.error("[Worker] Task %d hilang dari DB setelah transkripsi.", task_id)
            return

        if cancelled:
            update_task_status(
                db, task,
                status        = "cancelled",
                progress      = 0,
                error_message = "Task dibatalkan oleh admin.",
                partial_result= "",      # bersihkan hasil sementara
                segments_json = "[]",
            )
            logger.info("[Worker] Task %d tersimpan sebagai 'cancelled'.", task_id)

        elif output is not None:
            raw_text = output.get("text", "") or ""

            # Terapkan court mode pada hasil final jika aktif
            if is_court_mode:
                raw_text = apply_court_mode(raw_text, is_court_mode=True)

            encrypted_result = encrypt_text(raw_text) if raw_text else ""

            update_task_status(
                db, task,
                status          = "completed",
                progress        = 100,
                result          = encrypted_result,
                processing_time = output.get("processing_time"),
                audio_duration  = output.get("audio_duration"),
                partial_result  = "",        # bersihkan partial setelah final tersimpan
                segments_json   = "[]",
            )
            logger.info("[Worker] Task %d tersimpan sebagai 'completed' (encrypted).", task_id)

        else:
            update_task_status(
                db, task,
                status        = "failed",
                progress      = 0,
                error_message = error_message or "Terjadi kesalahan tidak diketahui.",
                partial_result= "",
                segments_json = "[]",
            )
            logger.warning("[Worker] Task %d tersimpan sebagai 'failed'.", task_id)


def _only_update_progress(task_id: int, progress: int) -> None:
    """
    Update HANYA kolom progress — tidak mengubah status sama sekali.

    FIX BUG: Versi lama selalu mengirim status="processing" sehingga bisa
    menimpa status "completed" atau "failed" jika ada race condition.
    Sekarang kita update progress via SQL langsung tanpa menyentuh kolom status.
    """
    try:
        with get_db_context() as db:
            db.query(TranskripsiTask).filter(
                TranskripsiTask.id == task_id,
                # Guard: hanya update jika masih processing, jangan overwrite status final
                TranskripsiTask.status == "processing",
            ).update(
                {"progress": progress},
                synchronize_session=False,
            )
    except Exception as exc:
        logger.debug(
            "[Worker] Gagal update progress task %d ke %d: %s",
            task_id, progress, exc,
        )


def _flush_segments(task_id: int, segs: list[dict], is_court_mode: bool) -> None:
    """
    Flush segmen terkumpul ke DB sebagai partial_result + segments_json.
    Dipanggil dari callback setiap _SEGMENT_FLUSH_EVERY segmen.
    partial_result adalah gabungan teks semua segmen sejauh ini (plaintext).
    """
    try:
        partial_text = " ".join(s["text"] for s in segs if s.get("text"))
        if is_court_mode:
            partial_text = apply_court_mode(partial_text, is_court_mode=True)

        segs_json_str = json.dumps(segs, ensure_ascii=False)

        # Hitung progress berdasarkan durasi: segmen terakhir / estimasi total
        # Kita pakai 40–80% range untuk transkripsi (40=model loaded, 80=hampir selesai)
        last_end    = segs[-1]["end"] if segs else 0
        # progress estimatif — akan di-overwrite saat selesai dengan 100%
        est_progress = min(79, 40 + int(last_end / max(last_end + 1, 1) * 39))

        with get_db_context() as db:
            db.query(TranskripsiTask).filter(
                TranskripsiTask.id     == task_id,
                TranskripsiTask.status == "processing",   # guard: jangan overwrite status final
            ).update(
                {
                    "partial_result": partial_text,
                    "segments_json":  segs_json_str,
                    "progress":       est_progress,
                },
                synchronize_session=False,
            )
    except Exception as exc:
        logger.debug("[Worker] Gagal flush segmen task %d: %s", task_id, exc)


def _set_cancelled(task_id: int) -> None:
    """Set status cancelled untuk task yang belum sempat mulai."""
    try:
        with get_db_context() as db:
            db.query(TranskripsiTask).filter(
                TranskripsiTask.id == task_id,
            ).update(
                {"status": "cancelled", "progress": 0},
                synchronize_session=False,
            )
    except Exception as exc:
        logger.debug("[Worker] Gagal set cancelled task %d: %s", task_id, exc)