"""Task 02 - Process flow diagrams.

Reads ./out/config/lab_config.yaml and writes the Mermaid documentation of
the laboratory workflow: ./out/report/process_flow.md plus one standalone
./out/report/diagrams/D<n>.mermaid file per diagram (D1-D6), and the
./out/.done/02.ok completion marker. Spec: docs/specs/02_process_flow_diagrams.md.
"""

from __future__ import annotations

import logging
import sys
from pathlib import Path

from common.pipeline import (
	check_dependencies,
	load_lab_config,
	parse_task_args,
	remove_stale_marker,
	write_marker,
)

logger = logging.getLogger(__name__)

TASK_ID = "02"
DEPENDS_ON = ["01"]

DIAGRAM_TYPES = ("flowchart", "sequenceDiagram", "erDiagram")


def d1_coupon_assembly(cfg: dict) -> str:
	sub = cfg["material_system"]["substrate"]
	dims = "x".join(str(d) for d in sub["dimensions_mm"])
	return f"""flowchart LR
    A([{sub['material']} Base {dims} mm]) --> B[Cleaning]
    B --> B1([Rinsed in deionized water])
    B1 --> B2([Sonicated in cleaning solution])
    B --> C[Smearing Adhesive]
    C --> C1([{cfg['material_system']['adhesion_layer']['material']} adhesive coating])
    C --> D[Vapor Deposition System<br/>Kurt J. Lesker PVD 200]
    D --> D1([PtAu sputter coating, gradient])
    D --> E([Test Coupon])
"""


def d2_characterization_batch_assembly(cfg: dict) -> str:
	h = cfg["hierarchy"]
	t_lo, t_hi = cfg["material_system"]["coating"]["thickness_um_range"]
	lines = [
		"flowchart LR",
		"    TC([Test Coupon]) --> S[SIMTRA simulation<br/>composition / atom energies]",
		f"    S --> OP[Optical profilometry<br/>film thickness {t_lo}-{t_hi} um]",
		f"    OP --> TW[Assemble Test Wafer<br/>x{h['coupons_per_wafer']} coupons]",
		f"    TW --> TB[Assemble Test Batch<br/>x{h['wafers_per_batch']} wafers]",
	]
	for row in cfg["deposition_matrix"]:
		code = row["batch_code"]
		lines.append(
			f"    TB --> {code}[{code}<br/>Pt {row['pt_power_W']} W / Au {row['au_power_W']} W]"
		)
	return "\n".join(lines) + "\n"


def d3_testing_tree(cfg: dict) -> str:
	v = cfg["volumes"]
	xrf = "x".join(str(g) for g in v["xrf_grid"])
	return f"""flowchart TD
    TC([Test Coupon]) --> F[Friction<br/>RAPID 6-probe tribometer]
    F --> FA([Lab Air -> COF dataset mu_normal])
    F --> FN([Dry N2 -> COF dataset mu_dry_nit])
    TC --> NI[Nanoindentation<br/>Bruker TI980]
    NI --> NIR([Hardness, reduced modulus<br/>{v['nanoindentation_indents_per_coupon']} indents per coupon])
    TC --> AF[AFM]
    AF --> AFR([Topography -> Ra / Rq])
    TC --> XR[micro-XRF<br/>Bruker M4 Tornado]
    XR --> XRR([Pt/Au wt% map, {xrf} grid])
"""


def d4_tribometer_session(cfg: dict) -> str:
	f = cfg["friction_assignment"]
	return f"""sequenceDiagram
    actor Op as Operator
    participant Prep as Sample Prep
    participant Plan as Excel Test Plan
    participant SW as Control Software
    participant TR as RAPID Tribometer

    Op->>Prep: Request test
    Op->>Prep: Get samples
    Op->>Prep: Load samples onto plates
    Op->>Prep: Load ball holders ({f['counterfaces_per_holder']} counterfaces per holder)
    Op->>Plan: Create test plan<br/>(Sample Plate, Plate Location, Sample ID, Save Location,<br/>Folder Name, X/Y/Z Offset, Load, Iterations)
    Op->>TR: Transfer plates to tribometer
    Op->>TR: Set up platter
    Op->>SW: Activate software
    Op->>SW: Load plan
    SW->>TR: Activate tribometer
    TR-->>SW: Run reciprocating tests (execution loop, see D5)
    Op->>SW: Software stop
    SW-->>Op: Save avg files
    Op->>TR: Remove plates and platters
    Op->>SW: Close software
    Op->>TR: Pull equipment out
"""


def d5_execution_loop(cfg: dict) -> str:
	f = cfg["friction_assignment"]
	tracks_per_run = f["coupons_per_run"] * f["tracks_per_friction_coupon"]
	return f"""flowchart TD
    L[Load all counterfaces] --> P{{{f['plates_per_run']} plates}}
    P -->|for each plate| PR{{{f['probes_per_plate']} probes in parallel}}
    PR -->|for each probe| C{{{f['coupons_per_probe_square']} coupons on coupon square}}
    C -->|for each coupon| T[Draw track]
    T --> RC[Rotate counterface - fresh ball per track]
    RC -->|{f['tracks_per_friction_coupon']} tracks per coupon| T
    T --> DONE([{f['coupons_per_run']} coupons x {f['tracks_per_friction_coupon']} tracks = {tracks_per_run} tracks per run])
"""


def d6_data_hierarchy(cfg: dict) -> str:
	return """erDiagram
    DEPOSITION_RUN ||--|| SIMTRA_PROFILE : produces
    DEPOSITION_RUN ||--|| BATCH : creates
    BATCH ||--|{ WAFER : contains
    WAFER ||--|{ COUPON : contains
    COUPON ||--|| XRF_MAP : has
    COUPON ||--|| NANOINDENTATION : has
    COUPON ||--|| AFM : has
    COUPON ||--|| PROFILOMETRY : has
    COUPON ||--o{ TRACK : "friction coupons only"
    TRACK ||--|{ CYCLE : has
    TRACK ||--|{ LOOP_POINT : has
    TRACK ||--|| WEAR : has
"""


def deposition_table(cfg: dict) -> str:
	lines = [
		"| Batch | Pt:Au gun tilt | Pt power W | Au power W | Pt discharge V | Au discharge V |",
		"|---|---|---|---|---|---|",
	]
	for r in cfg["deposition_matrix"]:
		lines.append(
			f"| {r['batch_code']} | {r['pt_gun_tilt_deg']}:{r['au_gun_tilt_deg']} deg "
			f"| {r['pt_power_W']} | {r['au_power_W']} | {r['pt_discharge_V']} | {r['au_discharge_V']} |"
		)
	return "\n".join(lines) + "\n"


def build_diagrams(cfg: dict) -> list[tuple[str, str, str, str]]:
	"""Return (file_stem, title, description, mermaid_text) per diagram."""
	h = cfg["hierarchy"]
	return [
		(
			"D1",
			"Coupon Assembly",
			"Substrate preparation, adhesion layer, and gradient Pt-Au sputter deposition.",
			d1_coupon_assembly(cfg),
		),
		(
			"D2",
			"Characterization and Batch Assembly",
			f"Per-coupon simulation and thickness characterization, then assembly into wafers "
			f"(x{h['coupons_per_wafer']} coupons) and batches (x{h['wafers_per_batch']} wafers), "
			f"deposited per the parameter table below.",
			d2_characterization_batch_assembly(cfg),
		),
		(
			"D3",
			"Testing Tree",
			"The four characterization/testing paths every coupon can take.",
			d3_testing_tree(cfg),
		),
		(
			"D4",
			"Tribometer Session Sequence",
			"Operator workflow for one RAPID tribometer session, from test request to teardown.",
			d4_tribometer_session(cfg),
		),
		(
			"D5",
			"Execution Loop",
			"Nested iteration executed by the tribometer within one run.",
			d5_execution_loop(cfg),
		),
		(
			"D6",
			"Data Hierarchy",
			"Entity containment and per-entity measurement datasets (ERD-style).",
			d6_data_hierarchy(cfg),
		),
	]


def validate_diagram(stem: str, text: str) -> None:
	first = text.splitlines()[0].strip() if text.strip() else ""
	if not first.startswith(DIAGRAM_TYPES):
		raise ValueError(f"{stem}: unexpected diagram header {first!r}")


def write_outputs(cfg: dict, out_root: Path) -> tuple[Path, int, int]:
	diagrams_dir = out_root / "report" / "diagrams"
	diagrams_dir.mkdir(parents=True, exist_ok=True)
	diagrams = build_diagrams(cfg)

	md_parts = [
		"# Process Flow - Simulated Tribology Laboratory\n",
		"Generated by make_process_flow.py (task 02) from out/config/lab_config.yaml.",
		"Standalone diagram sources: ./diagrams/D1..D6.mermaid.\n",
	]
	total_diagram_bytes = 0
	for stem, title, description, text in diagrams:
		validate_diagram(stem, text)
		path = diagrams_dir / f"{stem}.mermaid"
		with open(path, "w", encoding="utf-8", newline="\n") as fh:
			fh.write(text)
		total_diagram_bytes += len(text.encode("utf-8"))
		logger.info("wrote %s (%s)", path, title)
		md_parts.append(f"## {stem} - {title}\n")
		md_parts.append(description + "\n")
		md_parts.append(f"```mermaid\n{text}```\n")
		if stem == "D2":
			md_parts.append("Deposition matrix:\n")
			md_parts.append(deposition_table(cfg))

	md_text = "\n".join(md_parts)
	md_path = out_root / "report" / "process_flow.md"
	with open(md_path, "w", encoding="utf-8", newline="\n") as fh:
		fh.write(md_text)
	return md_path, len(md_text.encode("utf-8")), total_diagram_bytes


def main() -> int:
	args = parse_task_args("Task 02: write process flow diagrams")
	out_root: Path = args.out_root
	try:
		check_dependencies(out_root, DEPENDS_ON)
		remove_stale_marker(out_root, TASK_ID)
		cfg = load_lab_config(out_root)
		md_path, md_bytes, diagram_bytes = write_outputs(cfg, out_root)
		marker_path = write_marker(out_root, TASK_ID, {
			"process_flow_file": md_path.as_posix(),
			"process_flow_bytes": md_bytes,
			"diagrams": 6,
			"diagram_bytes": diagram_bytes,
		})
	except Exception:
		logger.critical("task %s failed", TASK_ID, exc_info=True)
		return 1

	print(f"task 02 ok: {md_path} ({md_bytes} bytes)")
	print(f"diagrams: 6 files under {out_root / 'report' / 'diagrams'} ({diagram_bytes} bytes)")
	print(f"marker: {marker_path}")
	return 0


if __name__ == "__main__":
	sys.exit(main())