#!/usr/bin/env python3
"""记录 MAOMOMO 文章配图任务状态。

本脚本用于把图片生成进度写入状态文件，避免只用聊天记录声称已完成。
"""

from __future__ import annotations

import argparse
from datetime import datetime, timezone
import json
from pathlib import Path
import shutil
import sys
from typing import Any, Dict, Iterable, List, Optional, Tuple


JOBS_FILE = "slide_jobs.json"
STATE_FILE = "slide_run_state.json"


def now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def die(message: str, code: int = 1) -> None:
    print(f"错误：{message}", file=sys.stderr)
    raise SystemExit(code)


def read_text(path: Path) -> str:
    return path.read_text(encoding="utf-8")


def write_json(path: Path, data: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8", newline="\n") as handle:
        json.dump(data, handle, ensure_ascii=False, indent=2)
        handle.write("\n")


def read_json(path: Path) -> Any:
    try:
        return json.loads(read_text(path))
    except FileNotFoundError:
        die(f"找不到文件：{path}")
    except json.JSONDecodeError as exc:
        die(f"JSON 格式错误：{path}: {exc}")


def project_path(project: str) -> Path:
    path = Path(project).expanduser().resolve()
    if not path.exists():
        die(f"项目目录不存在：{path}")
    if not path.is_dir():
        die(f"不是项目目录：{path}")
    return path


def resolve_in_project(project: Path, value: str) -> Path:
    path = Path(value)
    if not path.is_absolute():
        path = project / path
    return path.resolve()


def load_job_doc(project: Path) -> Dict[str, Any]:
    raw = read_json(project / JOBS_FILE)
    if isinstance(raw, list):
        return {"version": 1, "jobs": raw}
    if not isinstance(raw, dict):
        die(f"{JOBS_FILE} 必须是对象或数组。")
    jobs = raw.get("jobs")
    if not isinstance(jobs, list):
        die(f"{JOBS_FILE} 缺少 jobs 数组。")
    return raw


def load_state_doc(project: Path) -> Dict[str, Any]:
    path = project / STATE_FILE
    if not path.exists():
        return {"version": 1, "events": []}
    raw = read_json(path)
    if not isinstance(raw, dict):
        die(f"{STATE_FILE} 必须是对象。")
    if not isinstance(raw.get("events"), list):
        raw["events"] = []
    return raw


def save_docs(project: Path, job_doc: Dict[str, Any], state_doc: Dict[str, Any]) -> None:
    write_json(project / JOBS_FILE, job_doc)
    write_json(project / STATE_FILE, state_doc)


def job_key(job: Dict[str, Any]) -> Optional[str]:
    for key in ("id", "job_id", "slide", "name"):
        value = job.get(key)
        if isinstance(value, str) and value:
            return value
    return None


def find_job(job_doc: Dict[str, Any], wanted: str) -> Dict[str, Any]:
    for job in job_doc["jobs"]:
        if isinstance(job, dict) and job_key(job) == wanted:
            return job
    die(f"未找到图片 job：{wanted}")
    return {}


def event(state_doc: Dict[str, Any], kind: str, job_id: str, payload: Dict[str, Any]) -> None:
    state_doc.setdefault("events", []).append(
        {
            "time": now_iso(),
            "event": kind,
            "job": job_id,
            **payload,
        }
    )


def infer_origin_out(prompt_path: Path, data: Dict[str, Any]) -> str:
    for key in ("origin_out", "out", "output", "selected_output"):
        value = data.get(key)
        if isinstance(value, str) and value:
            return value
    file_name = data.get("file_name")
    if isinstance(file_name, str) and file_name:
        return str(Path("origin_image") / Path(file_name).name)
    return str(Path("origin_image") / f"{prompt_path.stem}.png")


def infer_asset_out(data: Dict[str, Any], origin_out: str) -> str:
    value = data.get("asset_out")
    if isinstance(value, str) and value:
        return value
    return str(Path("assets") / Path(origin_out).name)


def cmd_init(args: argparse.Namespace) -> int:
    project = project_path(args.project)
    prompts_dir = resolve_in_project(project, args.prompts_dir)
    if not prompts_dir.exists():
        die(f"prompts 目录不存在：{prompts_dir}")

    jobs: List[Dict[str, Any]] = []
    for prompt_path in sorted(prompts_dir.glob("*.json")):
        data = read_json(prompt_path)
        if not isinstance(data, dict):
            die(f"prompt job 必须是 JSON 对象：{prompt_path}")
        job_id = str(data.get("id") or data.get("job_id") or prompt_path.stem)
        origin_out = infer_origin_out(prompt_path, data)
        jobs.append(
            {
                "id": job_id,
                "prompt_file": str(prompt_path.relative_to(project)),
                "origin_out": origin_out,
                "asset_out": infer_asset_out(data, origin_out),
                "status": "pending",
                "selected_backend": args.selected_backend,
                "created_at": now_iso(),
            }
        )

    if not jobs:
        die(f"没有找到 prompt JSON：{prompts_dir}")

    job_doc = {"version": 1, "jobs": jobs}
    state_doc = {"version": 1, "events": [{"time": now_iso(), "event": "init", "job_count": len(jobs)}]}
    save_docs(project, job_doc, state_doc)
    print(json.dumps({"status": "initialized", "jobs": len(jobs)}, ensure_ascii=False, indent=2))
    return 0


def cmd_status(args: argparse.Namespace) -> int:
    project = project_path(args.project)
    job_doc = load_job_doc(project)
    counts: Dict[str, int] = {}
    for job in job_doc["jobs"]:
        if not isinstance(job, dict):
            continue
        status = str(job.get("status") or "unknown")
        counts[status] = counts.get(status, 0) + 1

    print(json.dumps({"project": str(project), "counts": counts}, ensure_ascii=False, indent=2))
    for job in job_doc["jobs"]:
        if isinstance(job, dict):
            print(f"{job_key(job) or '<unknown>'}\t{job.get('status', 'unknown')}\t{job.get('origin_out', '')}")
    return 0


def cmd_dispatch(args: argparse.Namespace) -> int:
    project = project_path(args.project)
    job_doc = load_job_doc(project)
    state_doc = load_state_doc(project)
    job = find_job(job_doc, args.job)
    prompt_file = args.prompt_file or job.get("prompt_file")
    if not isinstance(prompt_file, str) or not prompt_file:
        die("缺少 prompt 文件路径。")

    job["status"] = "dispatched"
    job["dispatch"] = {
        "agent_id": args.agent_id,
        "agent_nickname": args.agent_nickname,
        "prompt_file": prompt_file,
        "time": now_iso(),
    }
    event(state_doc, "dispatch", args.job, job["dispatch"])
    save_docs(project, job_doc, state_doc)
    print(json.dumps({"status": "dispatched", "job": args.job}, ensure_ascii=False, indent=2))
    return 0


def copy_if_needed(source: Path, target: Path) -> None:
    if not source.exists():
        die(f"结果图片不存在：{source}")
    target.parent.mkdir(parents=True, exist_ok=True)
    try:
        if source.resolve() == target.resolve():
            return
    except FileNotFoundError:
        pass
    shutil.copyfile(source, target)


def cmd_result(args: argparse.Namespace) -> int:
    project = project_path(args.project)
    job_doc = load_job_doc(project)
    state_doc = load_state_doc(project)
    job = find_job(job_doc, args.job)
    source = resolve_in_project(project, args.selected_source)

    origin_value = args.origin_out or job.get("origin_out")
    if not isinstance(origin_value, str) or not origin_value:
        die("缺少 origin_out。")
    origin_out = resolve_in_project(project, origin_value)
    copy_if_needed(source, origin_out)

    asset_value = args.asset_out or job.get("asset_out")
    if isinstance(asset_value, str) and asset_value:
        copy_if_needed(origin_out, resolve_in_project(project, asset_value))

    result = {
        "agent_id": args.agent_id,
        "backend_used": args.backend_used,
        "selected_source": str(source),
        "origin_out": str(origin_out.relative_to(project) if origin_out.is_relative_to(project) else origin_out),
        "asset_out": asset_value,
        "qa_note": args.qa_note,
        "time": now_iso(),
    }
    job["status"] = "recorded"
    job["result"] = result
    event(state_doc, "result", args.job, result)
    save_docs(project, job_doc, state_doc)
    print(json.dumps({"status": "recorded", "job": args.job, "origin_out": result["origin_out"]}, ensure_ascii=False, indent=2))
    return 0


def cmd_blocker(args: argparse.Namespace) -> int:
    project = project_path(args.project)
    job_doc = load_job_doc(project)
    state_doc = load_state_doc(project)
    job = find_job(job_doc, args.job)
    blocker = {
        "agent_id": args.agent_id,
        "reason": args.reason,
        "evidence": args.evidence,
        "time": now_iso(),
    }
    job["status"] = "blocked"
    job["blocker"] = blocker
    event(state_doc, "blocker", args.job, blocker)
    save_docs(project, job_doc, state_doc)
    print(json.dumps({"status": "blocked", "job": args.job, "reason": args.reason}, ensure_ascii=False, indent=2))
    return 0


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Record MAOMOMO article image job state.")
    subparsers = parser.add_subparsers(dest="command", required=True)

    init = subparsers.add_parser("init", help="Create slide_jobs.json from prompts/*.json.")
    init.add_argument("project")
    init.add_argument("--prompts-dir", default="prompts")
    init.add_argument("--selected-backend", required=True)
    init.set_defaults(func=cmd_init)

    status = subparsers.add_parser("status", help="Print job status summary.")
    status.add_argument("project")
    status.set_defaults(func=cmd_status)

    dispatch = subparsers.add_parser("dispatch", help="Record a dispatched image job.")
    dispatch.add_argument("project")
    dispatch.add_argument("--job", required=True)
    dispatch.add_argument("--agent-id", required=True)
    dispatch.add_argument("--agent-nickname", default="")
    dispatch.add_argument("--prompt-file")
    dispatch.set_defaults(func=cmd_dispatch)

    result = subparsers.add_parser("result", help="Record a generated image result.")
    result.add_argument("project")
    result.add_argument("--job", required=True)
    result.add_argument("--agent-id", required=True)
    result.add_argument("--backend-used", required=True)
    result.add_argument("--selected-source", required=True)
    result.add_argument("--origin-out")
    result.add_argument("--asset-out")
    result.add_argument("--qa-note", default="")
    result.set_defaults(func=cmd_result)

    blocker = subparsers.add_parser("blocker", help="Record a blocked image job.")
    blocker.add_argument("project")
    blocker.add_argument("--job", required=True)
    blocker.add_argument("--agent-id", required=True)
    blocker.add_argument("--reason", required=True)
    blocker.add_argument("--evidence", default="")
    blocker.set_defaults(func=cmd_blocker)
    return parser


def main(argv: Optional[Iterable[str]] = None) -> int:
    parser = build_parser()
    args = parser.parse_args(list(argv) if argv is not None else None)
    return args.func(args)


if __name__ == "__main__":
    raise SystemExit(main())