from __future__ import annotations

import hashlib
import json
import os
import shutil
import sqlite3
import threading
import time
import uuid
from pathlib import Path

from asrpipe import graph, lexicon, providers, refine, review, revise, vault
from asrpipe.config import CFG
from asrpipe.locking import acquire_recording_lock
from . import db
from .config import (
    DATA, DRAFTS, FAILED, INBOX, LOGS, MAX_ATTEMPTS,
    POLL_SECONDS, PROCESSED, PROCESSING, PUBLIC_URL, REVISION_QUEUE,
    STABLE_SECONDS, STATE, SUPPORTED, TRANSCRIPTS, ensure_dirs,
)
from .notify import send

_RECORDING_LOCK = None


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:
    """Backward-compatible helper for qwen3cam JSON responses."""
    return providers.transcript_from_qwen3cam_response(payload)


def transcribe(
    audio: Path,
    job_id: str,
    *,
    provider: str = "",
    model: str = "",
    title: str = "",
    user_context: str = "",
) -> tuple[Path, providers.TranscribeResult]:
    res = providers.transcribe_audio(
        audio,
        provider=provider or CFG.asr_provider,
        model=model,
        title=title,
        user_context=user_context,
    )
    response_path = TRANSCRIPTS / f"{job_id}.asr.json"
    try:
        response_path.write_text(
            json.dumps(res.raw_response, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )
    except Exception:
        pass
    transcript = TRANSCRIPTS / f"{job_id}.txt"
    transcript.write_text(res.transcript_text, encoding="utf-8")
    return transcript, res


def intake_header(payload: dict) -> str:
    blob = json.dumps(payload, ensure_ascii=False).replace("--", "- -")
    return f"<!-- intake: {blob} -->"


def replace_intake(raw: str, header: str) -> str:
    if raw.startswith("<!-- intake:"):
        _first, _sep, rest = raw.partition("\n")
        return header + ("\n" + rest if _sep else "")
    return header + "\n" + raw


def beat() -> None:
    ensure_dirs()
    (STATE / "worker.heartbeat").write_text(str(time.time()))


def _beat_loop(stop: threading.Event) -> None:
    while not stop.wait(15):
        try:
            beat()
        except Exception:
            pass


def process_audio(source: Path, job: dict) -> None:
    job_id = job["id"]
    attempts = int(job.get("attempts") or 0) + 1
    prov = providers.normalize_provider_name(job.get("asr_provider") or CFG.asr_provider)
    model = str(job.get("asr_model") or "")
    digest = sha256(source)
    prov_label = providers.PROVIDERS.get(prov, {}).get("label", prov)
    db.update_job(
        job_id, "transcribing", f"正在转写 ({prov_label})",
        sha256=digest, attempts=attempts, error_code="",
        asr_provider=prov, asr_model=model,
    )
    transcript, res = transcribe(
        source, job_id,
        provider=prov, model=model,
        title=str(job.get("title") or ""),
        user_context=str(job.get("context") or ""),
    )
    intake = {
        "job_id": job_id,
        "title": job.get("title", ""),
        "context": job.get("context", ""),
        "source": job.get("source", ""),
        "asr_provider": res.provider,
        "asr_model": res.model,
    }
    raw = transcript.read_text(encoding="utf-8")
    raw = replace_intake(raw, intake_header(intake))
    transcript.write_text(raw, encoding="utf-8")
    db.update_job(
        job_id, "refining", f"转写完成 ({res.provider}:{res.model})，正在生成整理稿",
        transcript=str(transcript), asr_provider=res.provider, asr_model=res.model,
    )
    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 rerefine_job(job: dict) -> None:
    """Re-run LLM refinement on an existing raw transcript without re-running ASR."""
    job_id = job["id"]
    transcript = Path(str(job.get("transcript") or ""))
    if not transcript.is_file():
        raise FileNotFoundError(f"transcript file missing for job {job_id}: {transcript}")
    raw = transcript.read_text(encoding="utf-8", errors="ignore")
    intake = {
        "job_id": job_id,
        "title": job.get("title", ""),
        "context": job.get("context", ""),
        "source": job.get("source", ""),
        "asr_provider": job.get("asr_provider") or CFG.asr_provider,
        "asr_model": job.get("asr_model") or "",
    }
    raw = replace_intake(raw, intake_header(intake))
    transcript.write_text(raw, encoding="utf-8")

    db.update_job(job_id, "refining", "正在基于现有转写重新生成精修草稿")
    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)
    review.delete_pending_for_job(
        transcript_path=str(transcript),
        draft_path=str(job.get("draft") or draft),
        job_id=job_id,
    )
    questions = review.enqueue_draft(str(transcript), draft)
    db.update_job(
        job_id, "review", f"整理稿已重新生成，{questions} 个问题待确认",
        transcript=str(transcript), draft=str(draft),
    )


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",
            asr_provider=CFG.asr_provider,
        )
    target = PROCESSING / path.name
    if target.exists():
        target = PROCESSING / f"{path.stem}-{job['id'][:8]}{path.suffix}"
    path.replace(target)
    if target.name != job.get("stored_name"):
        db.update_job(job["id"], job["status"], job.get("detail") or "已认领", stored_name=target.name)
        job = db.get_job(job["id"]) or job
    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)
            is_empty_audio = isinstance(exc, providers.EmptyAudioError)
            retry = (attempts < MAX_ATTEMPTS) and not is_empty_audio
            db.update_job(
                job["id"], "queued" if retry else "failed",
                f"处理失败，等待重试 ({exc})" if retry else f"处理失败：{exc}",
                error_code=type(exc).__name__,
            )
            log(f"job={job['id']} failed type={type(exc).__name__} retry={retry} err={exc}")
            try:
                if source.exists():
                    folder = INBOX if retry else FAILED
                    target = folder / source.name
                    if target.exists():
                        target = folder / f"{source.stem}-{job['id'][:8]}{source.suffix}"
                    shutil.move(str(source), target)
                    if retry:
                        old = time.time() - max(STABLE_SECONDS, 60)
                        os.utime(target, (old, old))
            except Exception as move_exc:
                log(f"job={job['id']} audio left in processing after failure: {move_exc}")


def _defer_revision(marker: Path, payload: dict, exc: BaseException) -> None:
    attempts = int(payload.get("attempts") or 0) + 1
    payload["attempts"] = attempts
    payload["last_error"] = str(exc)[:300]
    job_id = str(payload.get("job_id") or "")
    if attempts >= MAX_ATTEMPTS:
        marker.unlink(missing_ok=True)
        if job_id:
            db.update_job(job_id, "failed", f"修订失败：{exc}", error_code=type(exc).__name__)
        log(f"revision marker={marker.name} gave up after {attempts}: {exc}")
        return
    tmp = marker.with_name(f".{marker.name}.{os.getpid()}.tmp")
    tmp.write_text(json.dumps(payload, ensure_ascii=False), encoding="utf-8")
    tmp.replace(marker)
    if job_id:
        db.update_job(
            job_id, "review",
            f"修订失败，将重试 ({attempts}/{MAX_ATTEMPTS})：{exc}",
            error_code=type(exc).__name__,
        )
    log(f"revision marker={marker.name} retry {attempts}/{MAX_ATTEMPTS}: {exc}")


def process_revisions() -> None:
    for marker in sorted(REVISION_QUEUE.glob("*.json")):
        if marker.name.startswith("."):
            continue
        try:
            payload = json.loads(marker.read_text(encoding="utf-8"))
        except Exception as exc:
            log(f"revision marker={marker.name} unreadable: {exc}")
            _defer_revision(marker, {}, exc)
            continue
        try:
            action = str(payload.get("action") or "revise")
            job_id = str(payload.get("job_id") or "")
            if action == "rerefine":
                job = db.get_job(job_id)
                if job:
                    rerefine_job(job)
                marker.unlink(missing_ok=True)
                continue
            draft = Path(payload["draft"])
            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:
            _defer_revision(marker, payload, exc)


def _park_audio(path: Path) -> Path:
    dest = INBOX / path.name
    if dest.exists():
        dest = INBOX / f"{path.stem}-recovered{path.suffix}"
    shutil.move(str(path), dest)
    old = time.time() - max(STABLE_SECONDS, 60)
    os.utime(dest, (old, old))
    return dest


def recover() -> None:
    """Return audio left in processing/ to the inbox and requeue its job."""
    if not PROCESSING.exists():
        return
    files = [p for p in PROCESSING.iterdir() if p.is_file()]
    jobs = db.pending_recovery()
    claimed: set[str] = set()
    for path in files:
        match = None
        for job in jobs:
            if job["id"] in claimed:
                continue
            stored = str(job.get("stored_name") or "")
            if path.name == stored or (job["id"][:8] and job["id"][:8] in path.name):
                match = job
                break
        dest = _park_audio(path)
        if match:
            claimed.add(match["id"])
            db.update_job(
                match["id"], "queued", "worker 重启后重新排队",
                stored_name=dest.name, error_code="",
            )
            log(f"recovered job={match['id']} audio={dest.name}")
        else:
            log(f"recovered orphan audio={dest.name}")


def sync_vault_knowledge() -> None:
    try:
        vault.pull()
        lexicon.sync_from_vault()
        graph.sync()
    except Exception as exc:
        log(f"vault knowledge sync skipped: {exc}")


def backup_databases() -> None:
    backup_dir = 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 run_loop() -> None:
    os.environ.setdefault("ASR_ROLE", "worker")
    os.environ.setdefault("ASR_REVISION_QUEUE_DIR", str(REVISION_QUEUE))
    ensure_dirs()
    global _RECORDING_LOCK
    lock_path = Path(os.getenv("ASR_CONSUMER_LOCK", str(STATE / "consumer.lock")))
    _RECORDING_LOCK = acquire_recording_lock(lock_path)
    recover()
    sync_vault_knowledge()
    beat()
    stop_beat = threading.Event()
    threading.Thread(target=_beat_loop, args=(stop_beat,), name="asr-heartbeat", daemon=True).start()
    last_backup = 0.0
    last_vault_sync = time.time()
    while True:
        try:
            scan_inbox()
        except Exception as exc:
            log(f"scan_inbox failed: {exc}")
        try:
            process_revisions()
        except Exception as exc:
            log(f"process_revisions failed: {exc}")
        now = time.time()
        if now - last_vault_sync >= 900:
            sync_vault_knowledge()
            last_vault_sync = now
        if now - last_backup >= 3600:
            try:
                backup_databases()
            except Exception as exc:
                log(f"backup failed: {exc}")
            last_backup = now
        beat()
        time.sleep(POLL_SECONDS)


def main() -> None:
    run_loop()


if __name__ == "__main__":
    main()