#!/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"
FINAL_MANIFEST_FILE = "final_manifest.json"
GITHUB_BATCHES_FILE = "github_batches.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 relative_path_string(project: Path, path: Path) -> str:
    try:
        return str(path.resolve().relative_to(project))
    except ValueError:
        return str(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 load_batches_doc(project: Path) -> Dict[str, Any]:
    path = project / GITHUB_BATCHES_FILE
    if not path.exists():
        return {"version": 1, "batches": []}
    raw = read_json(path)
    if not isinstance(raw, dict):
        die(f"{GITHUB_BATCHES_FILE} 必须是对象。")
    if not isinstance(raw.get("batches"), list):
        raw["batches"] = []
    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,
                "aspect_ratio": data.get("aspect_ratio", ""),
                "visual_direction": data.get("visual_direction", ""),
                "delivery_batch": data.get("delivery_batch", ""),
                "qa_status": "pending",
                "entered_assets": False,
                "final": False,
                "github_pushed": False,
                "github_batch": "",
                "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):
            github = "github" if job.get("github_pushed") else "local"
            qa = job.get("qa_status", "")
            print(f"{job_key(job) or '<unknown>'}\t{job.get('status', 'unknown')}\t{qa}\t{github}\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")
    entered_assets = False
    if isinstance(asset_value, str) and asset_value:
        copy_if_needed(origin_out, resolve_in_project(project, asset_value))
        entered_assets = True

    result = {
        "agent_id": args.agent_id,
        "backend_used": args.backend_used,
        "selected_source": str(source),
        "origin_out": relative_path_string(project, origin_out),
        "asset_out": asset_value,
        "qa_status": "passed",
        "entered_assets": entered_assets,
        "qa_note": args.qa_note,
        "time": now_iso(),
    }
    job["status"] = "recorded"
    job["qa_status"] = "passed"
    job["entered_assets"] = entered_assets
    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 cmd_final_manifest(args: argparse.Namespace) -> int:
    project = project_path(args.project)
    job_doc = load_job_doc(project)
    state_doc = load_state_doc(project)
    assets: List[Dict[str, Any]] = []
    incomplete: List[str] = []

    for job in job_doc["jobs"]:
        if not isinstance(job, dict):
            continue
        key = job_key(job) or "<unknown>"
        if job.get("status") != "recorded":
            incomplete.append(f"{key}: status={job.get('status', 'unknown')}")
            continue
        if job.get("qa_status") != "passed":
            incomplete.append(f"{key}: qa_status={job.get('qa_status', 'unknown')}")
            continue
        if job.get("entered_assets") is not True:
            incomplete.append(f"{key}: entered_assets=false")
            continue

        result = job.get("result") if isinstance(job.get("result"), dict) else {}
        asset_value = result.get("asset_out") or job.get("asset_out")
        if not isinstance(asset_value, str) or not asset_value:
            incomplete.append(f"{key}: 缺少 asset_out")
            continue

        asset_path = resolve_in_project(project, asset_value)
        if not asset_path.exists():
            incomplete.append(f"{key}: assets 文件不存在：{asset_value}")
            continue

        job["final"] = True
        assets.append(
            {
                "id": key,
                "prompt_file": job.get("prompt_file", ""),
                "origin_out": result.get("origin_out") or job.get("origin_out", ""),
                "asset_out": relative_path_string(project, asset_path),
                "qa_status": job.get("qa_status", ""),
                "qa_note": result.get("qa_note", ""),
                "aspect_ratio": job.get("aspect_ratio", ""),
                "visual_direction": job.get("visual_direction", ""),
                "delivery_batch": job.get("delivery_batch", ""),
                "github_pushed": bool(job.get("github_pushed")),
                "github_batch": job.get("github_batch", ""),
            }
        )

    if incomplete and not args.allow_incomplete:
        die("仍有未进入最终 manifest 的图片 job：\n" + "\n".join(incomplete))
    if not assets:
        die("没有可写入 final_manifest.json 的最终图片。")

    manifest_path = resolve_in_project(project, args.out)
    manifest = {
        "version": 1,
        "generated_at": now_iso(),
        "project": project.name,
        "asset_count": len(assets),
        "assets": assets,
    }
    write_json(manifest_path, manifest)
    event(
        state_doc,
        "final_manifest",
        "*",
        {
            "path": relative_path_string(project, manifest_path),
            "asset_count": len(assets),
            "allow_incomplete": args.allow_incomplete,
        },
    )
    save_docs(project, job_doc, state_doc)
    print(
        json.dumps(
            {
                "status": "final_manifest_written",
                "path": relative_path_string(project, manifest_path),
                "assets": len(assets),
            },
            ensure_ascii=False,
            indent=2,
        )
    )
    return 0


def cmd_github_batch(args: argparse.Namespace) -> int:
    project = project_path(args.project)
    job_doc = load_job_doc(project)
    state_doc = load_state_doc(project)
    batches_doc = load_batches_doc(project)

    rel_paths: List[str] = []
    for value in args.paths:
        path = resolve_in_project(project, value)
        if not path.exists():
            die(f"批次文件不存在：{value}")
        rel_paths.append(relative_path_string(project, path))

    path_set = set(rel_paths)
    pushed_jobs: List[str] = []
    for job in job_doc["jobs"]:
        if not isinstance(job, dict):
            continue
        result = job.get("result") if isinstance(job.get("result"), dict) else {}
        candidate_values = [job.get("asset_out"), result.get("asset_out")]
        normalized = {
            relative_path_string(project, resolve_in_project(project, value))
            for value in candidate_values
            if isinstance(value, str) and value
        }
        if normalized.intersection(path_set):
            job["github_pushed"] = True
            job["github_batch"] = args.batch
            if isinstance(job.get("result"), dict):
                job["result"]["github_pushed"] = True
                job["result"]["github_batch"] = args.batch
            pushed_jobs.append(job_key(job) or "<unknown>")

    batch_record = {
        "batch": args.batch,
        "commit": args.commit,
        "url": args.url,
        "release": args.release,
        "paths": rel_paths,
        "pushed_jobs": pushed_jobs,
        "note": args.note,
        "time": now_iso(),
    }
    batches = [item for item in batches_doc["batches"] if not (isinstance(item, dict) and item.get("batch") == args.batch)]
    batches.append(batch_record)
    batches_doc["batches"] = batches

    manifest_path = project / FINAL_MANIFEST_FILE
    if manifest_path.exists():
        manifest = read_json(manifest_path)
        if isinstance(manifest, dict) and isinstance(manifest.get("assets"), list):
            for item in manifest["assets"]:
                if not isinstance(item, dict):
                    continue
                asset_value = item.get("asset_out")
                if not isinstance(asset_value, str) or not asset_value:
                    continue
                normalized = relative_path_string(project, resolve_in_project(project, asset_value))
                if normalized in path_set:
                    item["github_pushed"] = True
                    item["github_batch"] = args.batch
                    item["github_commit"] = args.commit
                    item["github_url"] = args.url
            write_json(manifest_path, manifest)

    write_json(project / GITHUB_BATCHES_FILE, batches_doc)
    event(state_doc, "github_batch", "*", batch_record)
    save_docs(project, job_doc, state_doc)
    print(
        json.dumps(
            {
                "status": "github_batch_recorded",
                "batch": args.batch,
                "paths": len(rel_paths),
                "pushed_jobs": pushed_jobs,
            },
            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)

    final_manifest = subparsers.add_parser("final-manifest", help="Write final_manifest.json from recorded assets.")
    final_manifest.add_argument("project")
    final_manifest.add_argument("--out", default=FINAL_MANIFEST_FILE)
    final_manifest.add_argument("--allow-incomplete", action="store_true")
    final_manifest.set_defaults(func=cmd_final_manifest)

    github_batch = subparsers.add_parser("github-batch", help="Record a GitHub delivery batch.")
    github_batch.add_argument("project")
    github_batch.add_argument("--batch", required=True)
    github_batch.add_argument("--commit", required=True)
    github_batch.add_argument("--url", required=True)
    github_batch.add_argument("--release", default="")
    github_batch.add_argument("--paths", nargs="+", required=True)
    github_batch.add_argument("--note", default="")
    github_batch.set_defaults(func=cmd_github_batch)
    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())