from __future__ import annotations

import fcntl
import hashlib
import json
import os
import shutil
import sqlite3
import subprocess
import time
import uuid
from pathlib import Path

from asrpipe import refine, review, revise, vault
from . import db
from .config import (DIARIZE_URL, DRAFTS, FAILED, INBOX, MAX_ATTEMPTS, META,
                     POLL_SECONDS, PROCESSED, PROCESSING, PUBLIC_URL,
                     REVISION_QUEUE, STABLE_SECONDS, STATE, SUPPORTED, TRANSCRIPTS, LOGS,
                     ensure_dirs)
from .notify import send


def log(message: str) -> None:
    ensure_dirs()
    with (LOGS / "worker.log").open("a", encoding="utf-8") as stream:
        stream.write(time.strftime("[%Y-%m-%d %H:%M:%S] ") + message + "\n")


def sha256(path: Path) -> str:
    h = hashlib.sha256()
    with path.open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()


def transcript_from_response(payload: dict) -> str:
    segments = payload.get("segments")
    if not isinstance(segments, list) or not segments:
        raise RuntimeError("ASR response contains no segments")
    vp = payload.get("voiceprint") or {}
    sid = str(vp.get("session_id") or "")
    matched = [f"{m.get('speaker_local')}={m.get('person_name')}" for m in vp.get("matches", [])
               if m.get("status") == "matched"]
    speakers = {str(s.get("speaker")) for s in segments if s.get("speaker") is not None}
    lines = [f"<!-- voiceprint_session: {sid} clusters: {len(speakers)} -->",
             "# 转写(Qwen3-ASR 1.7B + CAM++ 分角色 + 声纹)"
             + (" | 声纹命名: " + ", ".join(matched) if matched else ""), ""]
    turns: list[list[str]] = []
    for segment in segments:
        speaker, text = str(segment.get("speaker", "?")), str(segment.get("text") or "").strip()
        if not text:
            continue
        if turns and turns[-1][0] == speaker:
            turns[-1][1] += text
        else:
            turns.append([speaker, text])
    lines.extend(f"【{speaker}】{text}\n" for speaker, text in turns)
    return "\n".join(lines).strip() + "\n"


def transcribe(audio: Path, job_id: str) -> Path:
    response_path = TRANSCRIPTS / f"{job_id}.asr.json"
    result = subprocess.run([
        "curl", "-sS", "--fail-with-body", "--max-time", "900",
        "-F", f"file=@{audio}", "-F", "language=Chinese",
        "-F", "save_voiceprint_session=true", DIARIZE_URL,
    ], capture_output=True, timeout=930)
    if result.returncode:
        raise RuntimeError(f"ASR HTTP failed ({result.returncode})")
    response_path.write_bytes(result.stdout)
    try:
        payload = json.loads(result.stdout)
    except ValueError as exc:
        raise RuntimeError("ASR returned invalid JSON") from exc
    transcript = TRANSCRIPTS / f"{job_id}.txt"
    transcript.write_text(transcript_from_response(payload), encoding="utf-8")
    return transcript


def process_audio(source: Path, job: dict) -> None:
    job_id = job["id"]
    attempts = int(job.get("attempts") or 0) + 1
    digest = sha256(source)
    db.update_job(job_id, "transcribing", "正在进行分角色转写", sha256=digest,
                  attempts=attempts, error_code="")
    transcript = transcribe(source, job_id)
    intake = {"job_id": job_id, "title": job.get("title", ""),
              "context": job.get("context", ""), "source": job.get("source", "")}
    raw = transcript.read_text(encoding="utf-8")
    raw = f"<!-- intake: {json.dumps(intake, ensure_ascii=False)} -->\n" + raw
    transcript.write_text(raw, encoding="utf-8")
    db.update_job(job_id, "refining", "转写完成，正在生成整理稿", transcript=str(transcript))
    result = refine.refine(raw, date=time.strftime("%Y-%m-%d"),
                           transcript_path=str(transcript), queue_review=True)
    draft = DRAFTS / f"{job_id}.md"
    vault.write_draft_file(draft, result.draft)
    questions = review.enqueue_draft(str(transcript), draft)
    db.update_job(job_id, "review", f"整理稿已生成，{questions} 个问题待确认",
                  transcript=str(transcript), draft=str(draft))
    target = PROCESSED / source.name
    if target.exists():
        target = PROCESSED / f"{source.stem}-{job_id[:8]}{source.suffix}"
    shutil.move(str(source), target)
    send("ASR 草稿待确认", job.get("title") or job["original_name"],
         f"{PUBLIC_URL}/jobs/{job_id}/review/")


def claim(path: Path) -> tuple[Path, dict]:
    job = db.find_by_stored(path.name)
    if not job:
        job_id = uuid.uuid4().hex[:24]
        job = db.create_job(job_id, path.name, path.name, title=path.stem,
                            size=path.stat().st_size, source="ssh")
    target = PROCESSING / path.name
    if target.exists():
        target = PROCESSING / f"{path.stem}-{job['id'][:8]}{path.suffix}"
    path.replace(target)
    return target, job


def scan_inbox() -> None:
    now = time.time()
    for path in sorted(INBOX.iterdir()):
        if not path.is_file() or path.suffix.lower() not in SUPPORTED:
            continue
        if now - path.stat().st_mtime < STABLE_SECONDS:
            continue
        source, job = claim(path)
        try:
            process_audio(source, job)
        except Exception as exc:
            attempts = int((db.get_job(job["id"]) or {}).get("attempts") or 0)
            retry = attempts < MAX_ATTEMPTS
            db.update_job(job["id"], "queued" if retry else "failed",
                          "处理失败，等待重试" if retry else "处理失败",
                          error_code=type(exc).__name__)
            log(f"job={job['id']} failed type={type(exc).__name__}")
            target = INBOX / source.name if retry else FAILED / source.name
            if source.exists():
                shutil.move(str(source), target)


def process_revisions() -> None:
    for marker in sorted(REVISION_QUEUE.glob("*.json")):
        try:
            payload = json.loads(marker.read_text(encoding="utf-8"))
            draft = Path(payload["draft"])
            job_id = str(payload.get("job_id") or "")
            db.update_job(job_id, "publishing", "正在应用人工确认")
            revise.revise(draft)
            current = db.get_job(job_id) or {}
            if current.get("status") == "complete":
                send("ASR 正式记录已生成", current.get("title") or current.get("original_name", ""),
                     f"{PUBLIC_URL}/jobs/{job_id}/")
            elif current.get("status") != "review":
                db.update_job(job_id, "review", "定稿预览已生成，等待发布确认", draft=str(draft))
            marker.unlink(missing_ok=True)
        except Exception as exc:
            log(f"revision marker={marker.name} failed type={type(exc).__name__}")


def recover() -> None:
    for job in db.pending_recovery():
        if job["status"] in {"transcribing", "refining"}:
            candidate = next((p for p in PROCESSING.iterdir() if p.name == job["stored_name"]), None)
            if candidate:
                shutil.move(str(candidate), INBOX / candidate.name)
                db.update_job(job["id"], "queued", "worker 重启后重新排队")


def backup_databases() -> None:
    # STATE is local disk; durable copies belong on the NFS data volume.
    backup_dir = Path("/data/backups")
    backup_dir.mkdir(parents=True, exist_ok=True)
    for name in ("intake.db", "lexicon.db"):
        source = STATE / name
        if not source.exists():
            continue
        target = backup_dir / f"{name}.latest"
        temp = backup_dir / f".{name}.{os.getpid()}.tmp"
        with sqlite3.connect(source) as src, sqlite3.connect(temp) as dst:
            src.backup(dst)
        temp.replace(target)


def main() -> None:
    ensure_dirs()
    lock = (STATE / "worker.lock").open("w")
    fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
    recover()
    (STATE / "worker.heartbeat").write_text(str(time.time()))
    last_backup = 0.0
    while True:
        scan_inbox()
        process_revisions()
        if time.time() - last_backup >= 3600:
            backup_databases()
            last_backup = time.time()
        (STATE / "worker.heartbeat").write_text(str(time.time()))
        time.sleep(POLL_SECONDS)


if __name__ == "__main__":
    main()