"""DSH content-production worker.

The workbench API owns configuration, validation, queuing, human review and
delivery. This process leases ``campaign-*`` jobs from the shared SQLite queue
and executes the existing content generation pipeline.

Run:
    YAWAN_EXECUTION_BACKEND=dsh python -m app.dsh_worker

The worker must share the same YAWAN_DATA_DIR as the API process (SQLite,
knowledge bases, vector indexes and artifacts directory).
"""
from __future__ import annotations

import logging
import os
import signal
import socket
import threading
import time
import uuid

from sqlmodel import Session, select

from .agents import runner
from . import campaign_jobs
from . import planner
from .db import engine, init_db
from .models import (
    Campaign,
    CampaignRun,
    AgentSession,
    Notification,
    RunEvent,
    Segment,
    TaskQueueItem,
    now,
)
from . import progress
from . import task_queue

logger = logging.getLogger("yawen.dsh")

STOP = threading.Event()
WORKER_ID = os.getenv("YAWAN_DSH_WORKER_ID") or f"dsh-{socket.gethostname()}-{os.getpid()}-{uuid.uuid4().hex[:6]}"
POLL_SECONDS = float(os.getenv("YAWAN_DSH_POLL_SECONDS", "2.0"))
HEARTBEAT_SECONDS = int(os.getenv("YAWAN_DSH_HEARTBEAT_SECONDS", "60"))


def _record_event(session: Session, run: CampaignRun, event_type: str,
                  status: str = "", payload: dict | None = None) -> None:
    session.add(RunEvent(
        run_id=run.id, campaign_id=run.campaign_id, trace_id=run.trace_id,
        event_type=event_type, status=status, payload=payload or {},
    ))


def _set_campaign_review(session: Session, campaign_id: int, has_failure: bool | None = None) -> None:
    campaign = session.get(Campaign, campaign_id)
    if not campaign or campaign.status == "done":
        return
    segments = session.exec(select(Segment).where(Segment.campaign_id == campaign_id)).all()
    campaign.has_failure = has_failure if has_failure is not None else any(x.status == "failed" for x in segments)
    if all(x.status in ("approved", "skipped") for x in segments):
        campaign.status = "composing_ready"
    else:
        campaign.status = "human_review"
    session.add(campaign)


def _recover_orphan_run(session: Session, run: CampaignRun, payload: dict) -> None:
    """Reset stale runtime state before the leased worker owns the run.

    A lease is only stolen after heartbeat timeout, so any running session/segment
    here belongs to a dead worker. The same applies to an immediate in-process
    retry: the previous job thread is already gone.
    """
    for runtime_session in session.exec(select(AgentSession).where(
            AgentSession.run_id == run.id,
            AgentSession.status == "running")).all():
        runtime_session.status = "interrupted"
        runtime_session.updated_at = now()
        session.add(runtime_session)

    segment_id = int(payload.get("segment_id") or 0)
    if segment_id:
        seg = session.get(Segment, segment_id)
        if seg and seg.campaign_id == run.campaign_id and seg.status == "running":
            seg.status = "pending"
            session.add(seg)
    else:
        # 全量任务的上一次尝试可能已把 run 置为 failed；只要队列重新租约，
        # 所有遗留 running 段都属于上一个已结束 worker，应恢复为可执行状态。
        for seg in session.exec(select(Segment).where(
                Segment.campaign_id == run.campaign_id,
                Segment.status == "running")).all():
            seg.status = "pending"
            session.add(seg)
    session.commit()


def _run_full(session: Session, run: CampaignRun) -> str:
    campaign = session.get(Campaign, run.campaign_id)
    if run.cancel_requested:
        run.status = "cancelled"
        run.finished_at = now()
        if campaign:
            campaign.status = "human_review"
            session.add(campaign)
        session.add(run)
        _record_event(session, run, "run_finished", "cancelled", {"reason": "cancel_before_start"})
        session.commit()
        return run.status
    if run.pause_requested:
        run.status = "paused"
        if campaign:
            campaign.status = "human_review"
            session.add(campaign)
        session.add(run)
        _record_event(session, run, "run_finished", "paused", {"reason": "pause_before_start"})
        session.commit()
        return run.status
    run.status = "running"
    run.started_at = run.started_at or now()
    run.trace_id = run.trace_id or uuid.uuid4().hex
    run.checkpoint = {**(run.checkpoint or {}), "phase": "running", "worker_id": WORKER_ID}
    session.add(run)
    _record_event(session, run, "run_started", "running", {"worker_id": WORKER_ID})
    session.commit()

    result = runner.run_campaign(session, run.campaign_id, run_id=run.id)
    session.refresh(run)
    campaign = session.get(Campaign, run.campaign_id)

    if run.cancel_requested:
        run.status = "cancelled"
        run.finished_at = now()
        if campaign:
            campaign.status = "human_review"
    elif str(result.get("status")) == "paused" or run.pause_requested:
        run.status = "paused"
    elif str(result.get("status")) in ("failed", "error") or (campaign and campaign.has_failure):
        run.status = "failed"
        run.error = str(result.get("error") or "")[:500]
        run.finished_at = now()
    else:
        run.status = "awaiting_review"
        run.finished_at = now()

    run.checkpoint = {
        **(run.checkpoint or {}),
        "phase": "finished",
        "campaign_status": result.get("status"),
        "worker_id": WORKER_ID,
        "results": result.get("results", [])[-100:],
    }
    session.add(run)
    if campaign:
        session.add(campaign)
    _record_event(session, run, "run_finished", run.status, {"campaign_status": result.get("status")})
    session.commit()
    return run.status


def _run_step(session: Session, run: CampaignRun, segment_id: int) -> str:
    campaign = session.get(Campaign, run.campaign_id)
    if run.cancel_requested:
        run.status = "cancelled"
        run.finished_at = now()
        if campaign:
            campaign.status = "human_review"
            session.add(campaign)
        session.add(run)
        _record_event(session, run, "run_finished", "cancelled", {"reason": "cancel_before_start"})
        session.commit()
        return run.status
    if run.pause_requested:
        run.status = "paused"
        if campaign:
            campaign.status = "human_review"
            session.add(campaign)
        session.add(run)
        _record_event(session, run, "run_finished", "paused", {"reason": "pause_before_start"})
        session.commit()
        return run.status
    segment = session.get(Segment, segment_id)
    if not campaign or not segment or segment.campaign_id != campaign.id:
        raise ValueError(f"逐段任务无效：run={run.id}, segment={segment_id}")
    if segment.status not in ("pending", "rework", "failed"):
        run.status = "completed"
        run.error = ""
        run.finished_at = now()
        run.checkpoint = {**(run.checkpoint or {}), "phase": "skipped", "reason": segment.status}
        session.add(run)
        session.commit()
        return run.status

    run.status = "running"
    run.started_at = run.started_at or now()
    run.trace_id = run.trace_id or uuid.uuid4().hex
    run.checkpoint = {**(run.checkpoint or {}), "phase": "running", "mode": "step",
                      "segment_ids": [segment.id], "worker_id": WORKER_ID}
    session.add(run)
    _record_event(session, run, "run_started", "running", {"mode": "step", "segment_id": segment.id})
    session.commit()

    progress.publish(int(campaign.owner_user_id or 0), "run_step_start",
                     campaign_id=campaign.id, segment_id=segment.id,
                     title=segment.title, backend="dsh")
    try:
        result = runner.execute_segment(session, segment, campaign, run_id=run.id)
        if runner._supervisor_gate(session, segment, campaign, run.id) == "redo":
            segment = session.get(Segment, segment.id)
            result = runner.execute_segment(session, segment, campaign, run_id=run.id)
        status = str(result.get("status") or "failed")
        run.checkpoint = {**(run.checkpoint or {}), "last_result": status}
        if status == "waiting_approval":
            run.status = "awaiting_review"
        elif status in ("human_review", "rework", "approved"):
            run.status = "completed"
            run.finished_at = now()
        else:
            run.status = "failed"
            run.error = str(result.get("error") or "")[:500]
            run.finished_at = now()
        _set_campaign_review(session, campaign.id)
        session.add(run)
        _record_event(session, run, "run_finished", run.status, {"segment_status": status})
        session.commit()
        return run.status
    except Exception as exc:
        session.rollback()
        segment = session.get(Segment, segment_id)
        run = session.get(CampaignRun, run.id)
        campaign = session.get(Campaign, run.campaign_id)
        if segment:
            segment.status = "failed"
            segment.failure = {"type": "dsh_step_failed", "detail": str(exc)[:300]}
            session.add(segment)
        run.status = "failed"
        run.error = str(exc)[:500]
        run.finished_at = now()
        if campaign:
            campaign.status = "human_review"
            session.add(campaign)
        session.add(run)
        session.commit()
        raise


def _segment_redo(session: Session, payload: dict) -> None:
    segment_id = int(payload.get("segment_id") or 0)
    opinion = str(payload.get("opinion") or "")
    expected_agent = int(payload.get("expected_agent_id") or 0)
    expected_round = int(payload.get("expected_round") or 0)
    seg = session.get(Segment, segment_id)
    campaign = session.get(Campaign, seg.campaign_id) if seg else None
    if not seg or not campaign:
        raise ValueError("重做任务对应的内容段不存在")
    if seg.status == "running":
        # 队列租约过期后才可能重新领取到该确定性任务；旧 Worker 已被视为死亡。
        seg.status = "rework"
        session.add(seg)
        session.commit()
    if seg.status not in ("rework", "failed"):
        return
    if expected_agent and int(seg.agent_id or 0) != expected_agent:
        return
    if expected_round and int(seg.round or 0) != expected_round:
        return
    seg.status = "running"
    session.add(seg)
    session.commit()
    progress.publish(int(campaign.owner_user_id or 0), "redo_start", campaign_id=campaign.id,
                     segment_id=seg.id, round=seg.round, title=seg.title,
                     reason="user", backend="dsh")
    try:
        runner.execute_segment(session, seg, campaign, auto_retry_qc=False, qc_issues=opinion)
        seg2 = session.get(Segment, segment_id)
        progress.publish(int(campaign.owner_user_id or 0), "redo_done", campaign_id=campaign.id,
                         segment_id=segment_id, status=(seg2.status if seg2 else "?"),
                         round=(seg2.round if seg2 else 0), title=(seg2.title if seg2 else ""))
        _set_campaign_review(session, campaign.id)
        session.commit()
    except Exception as exc:
        session.rollback()
        seg = session.get(Segment, segment_id)
        campaign = session.get(Campaign, campaign.id)
        if seg:
            seg.status = "failed"
            seg.failure = {"type": "dsh_redo_failed", "detail": str(exc)[:300]}
            session.add(seg)
        if campaign:
            campaign.status = "human_review"
            campaign.has_failure = True
            session.add(campaign)
        session.commit()
        raise


def _segment_fix(session: Session, payload: dict) -> None:
    segment_id = int(payload.get("segment_id") or 0)
    seg = session.get(Segment, segment_id)
    if seg and seg.status in ("skipped", "approved"):
        # skip/disable 的管理动作已在 API 请求中完成；approved 说明人工已接管，旧任务作废。
        return
    if seg and seg.status == "running":
        # 上一执行进程租约超时，遗留 running 段恢复为 pending 后重试。
        seg.status = "pending"
        seg.failure = {}
        session.add(seg)
        session.commit()
    if seg and seg.status not in ("pending", "failed"):
        return
    fix = payload.get("fix") if isinstance(payload.get("fix"), dict) else {}
    try:
        result = runner.fix_segment(session, segment_id, fix, execute=True, apply_changes=False)
        if result.get("status") == "error":
            raise ValueError(str(result.get("error") or "修复执行失败"))
        seg = session.get(Segment, segment_id)
        if seg:
            _set_campaign_review(session, seg.campaign_id)
            session.commit()
    except Exception as exc:
        session.rollback()
        seg = session.get(Segment, segment_id)
        if seg:
            campaign = session.get(Campaign, seg.campaign_id)
            seg.status = "failed"
            seg.failure = {"type": "dsh_fix_failed", "detail": str(exc)[:300]}
            session.add(seg)
            if campaign:
                campaign.status = "human_review"
                campaign.has_failure = True
                session.add(campaign)
            session.commit()
        raise


def _campaign_parse(session: Session, payload: dict) -> None:
    campaign_id = int(payload.get("campaign_id") or 0)
    campaign_jobs.run_parse_job(session, campaign_id, str(payload.get("advice") or ""))


def _campaign_plan(session: Session, payload: dict) -> None:
    campaign_id = int(payload.get("campaign_id") or 0)
    template_id = int(payload.get("template_id") or 0)
    result = planner.run_plan_job(
        session,
        campaign_id,
        template_id=template_id,
        advice=str(payload.get("advice") or ""),
        created_by=int(payload.get("created_by") or 0),
    )
    if result.get("status") == "failed":
        logger.warning("DSH plan business failure: %s", str(result.get("error"))[:200])


def _campaign_compose(session: Session, payload: dict) -> None:
    campaign_id = int(payload.get("campaign_id") or 0)
    result = campaign_jobs.run_compose_job(session, campaign_id)
    if result.get("status") == "error":
        # 与 inline 同步接口一致：业务校验失败不做队列重试，工作台保留错误并允许用户重试点。
        logger.warning("DSH compose business failure: %s", str(result.get("error"))[:200])


def process_job(job: TaskQueueItem) -> None:
    with Session(engine) as session:
        payload = dict(job.payload or {})
        key = job.task_key or ""
        logger.info("DSH lease job %s: %s", job.id, key)

        if key.startswith("campaign-run:"):
            run_id = int(payload.get("run_id") or key.rsplit(":", 1)[-1])
            run = session.get(CampaignRun, run_id)
            if not run:
                raise ValueError(f"运行不存在：{key}")
            _recover_orphan_run(session, run, payload)
            session.refresh(run)
            try:
                if run.status in ("completed", "cancelled"):
                    return
                if payload.get("mode") == "step" or payload.get("segment_id"):
                    _run_step(session, run, int(payload.get("segment_id") or 0))
                else:
                    _run_full(session, run)
            finally:
                # 预算池是进程内对象；DSH 每次执行结束都必须释放，防止长驻 Worker 内存泄漏。
                from .agents.runtime import release_run_budget
                release_run_budget(run_id)
            return

        if key.startswith("campaign-segment-redo:"):
            _segment_redo(session, payload)
            return

        if key.startswith("campaign-segment-fix:"):
            _segment_fix(session, payload)
            return

        if key.startswith("campaign-plan:"):
            _campaign_plan(session, payload)
            return

        if key.startswith("campaign-parse:"):
            _campaign_parse(session, payload)
            return

        if key.startswith("campaign-compose:"):
            _campaign_compose(session, payload)
            return

        raise ValueError(f"DSH 不识别的内容生产任务：{key}")


def _mark_dead_job_state(session: Session, job: TaskQueueItem, error: str) -> None:
    """Keep business state visible when all queue retry attempts are exhausted."""
    payload = dict(job.payload or {})
    key = job.task_key or ""
    campaign_id = int(payload.get("campaign_id") or 0)
    run_id = int(payload.get("run_id") or 0)
    detail = str(error)[:300]

    if key.startswith("campaign-run:") or run_id:
        run = session.get(CampaignRun, run_id) if run_id else None
        if run:
            campaign_id = campaign_id or run.campaign_id
            run.status = "failed"
            run.error = detail
            run.finished_at = now()
            session.add(run)
    elif key.startswith("campaign-plan:") or key.startswith("campaign-parse:"):
        campaign = session.get(Campaign, campaign_id)
        if campaign:
            campaign.status = "plan_failed" if key.startswith("campaign-plan:") else "parse_failed"
            notes = list(campaign.notes or [])
            label = "规划" if key.startswith("campaign-plan:") else "拆解"
            notes.append(f"DSH {label}任务进入死信队列：{detail[:180]}")
            campaign.notes = notes[-20:]
            session.add(campaign)
    elif key.startswith("campaign-compose:"):
        campaign = session.get(Campaign, campaign_id)
        if campaign and campaign.status != "done":
            campaign.status = "composing_ready"
            session.add(campaign)

    segment_id = int(payload.get("segment_id") or 0)
    if segment_id:
        seg = session.get(Segment, segment_id)
        if seg and seg.status == "running":
            seg.status = "failed"
            seg.failure = {"type": "dsh_job_dead", "detail": detail}
            session.add(seg)
            campaign_id = campaign_id or seg.campaign_id

    campaign = session.get(Campaign, campaign_id) if campaign_id else None
    if campaign and campaign.status not in ("done", "parse_failed", "plan_failed", "composing_ready"):
        campaign.status = "human_review"
        campaign.has_failure = True
        session.add(campaign)
    session.commit()


def _heartbeat(job_id: int) -> threading.Thread:
    def loop():
        while not STOP.wait(HEARTBEAT_SECONDS):
            try:
                with Session(engine) as session:
                    ok = task_queue.heartbeat(session, TaskQueueItem, job_id, WORKER_ID)
                    if not ok:
                        return
            except Exception:
                logger.exception("DSH heartbeat failed: job=%s", job_id)
    thread = threading.Thread(target=loop, name=f"dsh-heartbeat-{job_id}", daemon=True)
    thread.start()
    return thread


def worker_loop() -> None:
    init_db()
    logger.info("DSH worker started: %s (poll=%ss, heartbeat=%ss)", WORKER_ID, POLL_SECONDS, HEARTBEAT_SECONDS)
    while not STOP.is_set():
        job = None
        try:
            with Session(engine) as session:
                job = task_queue.lease_next(session, TaskQueueItem, WORKER_ID,
                                            task_key_prefix="campaign-%")
            if job is None:
                STOP.wait(POLL_SECONDS)
                continue

            beat = _heartbeat(job.id)
            process_job(job)
            beat.join(timeout=1)
            with Session(engine) as session:
                task_queue.complete(session, TaskQueueItem, job.id, WORKER_ID)
            logger.info("DSH job completed: %s", job.id)
        except Exception as exc:
            logger.exception("DSH job failed: %s", getattr(job, "id", "?"))
            try:
                with Session(engine) as session:
                    if job is not None:
                        status = task_queue.fail(session, TaskQueueItem, job.id, WORKER_ID, str(exc))
                        if status == "dead":
                            _mark_dead_job_state(session, job, str(exc))
                            payload = dict(job.payload or {})
                            campaign_id = int(payload.get("campaign_id") or 0)
                            if not campaign_id and payload.get("run_id"):
                                run = session.get(CampaignRun, int(payload["run_id"]))
                                campaign_id = run.campaign_id if run else 0
                            if campaign_id:
                                session.add(Notification(
                                    type="failure", ref_id=campaign_id,
                                    message=f"DSH 内容生产任务进入死信队列：{str(exc)[:160]}",
                                ))
                                session.commit()
            except Exception:
                logger.exception("Failed to mark DSH job failed: %s", getattr(job, "id", "?"))
            STOP.wait(min(30.0, POLL_SECONDS * 3))
    logger.info("DSH worker stopped: %s", WORKER_ID)


def _signal_handler(signum, frame):  # pragma: no cover
    logger.info("DSH received signal %s, stopping...", signum)
    STOP.set()


def main() -> None:
    logging.basicConfig(
        level=os.getenv("YAWAN_DSH_LOG_LEVEL", "INFO").upper(),
        format="%(asctime)s %(levelname)s %(name)s: %(message)s",
    )
    if hasattr(signal, "SIGTERM"):
        signal.signal(signal.SIGTERM, _signal_handler)
    signal.signal(signal.SIGINT, _signal_handler)
    worker_loop()


if __name__ == "__main__":
    main()