"""Task 04 - Convert the CSV corpus to JSON-LD.

Reads ./out/csv/ (validated file-by-file against MANIFEST.csv) and writes
two ontology-annotated variants:
- FULL: one compact JSON-LD file per coupon under ./out/json/full/, bulk
  points embedded. Like spec 07's RDF modeling, bulk points (cycles, loop
  points, map points) carry no @type - the containing predicate types them;
  this keeps FULL lean, so its size lands below the 0.7-1.2 GB the spec
  sketched for a fully typed serialization.
- HYBRID: single ./out/json/hybrid/dataset.jsonld holding every metadata
  entity; every per-point array (cycles, loops, xrf, profilometry, indents)
  is replaced by a {sourceFile, rows, sha256} reference into the CSV tree.

Appends both sizes to ./out/bench/storage_sizes.csv and writes
./out/.done/04.ok. Spec: docs/specs/04_convert_json.md.
"""

from __future__ import annotations

import json
import logging
import shutil
import sys
from pathlib import Path

from rdflib import Graph

from common.corpus import CorpusReader, ValidationError
from common.pipeline import (
	check_dependencies,
	load_lab_config,
	parse_task_args,
	process_metrics,
	remove_stale_marker,
	write_marker,
)
from common.storage_sizes import update_storage_sizes

logger = logging.getLogger(__name__)

TASK_ID = "04"
DEPENDS_ON = ["01", "03"]

VOCAB = "https://sandia.gov/ontology/tribology#"
ID_BASE = "https://sandia.gov/ontology/tribology/id/"

CONTEXT = {
	"@vocab": VOCAB,
	"tribo": VOCAB,
	"qudt": "http://qudt.org/schema/qudt/",
	"unit": "http://qudt.org/vocab/unit/",
	"prov": "http://www.w3.org/ns/prov#",
	"batch": ID_BASE + "batch/",
	"wafer": ID_BASE + "wafer/",
	"coupon": ID_BASE + "coupon/",
	"track": ID_BASE + "track/",
	"run": ID_BASE + "run/",
	"instrument": ID_BASE + "instrument/",
	"partOf": {"@id": "tribo:partOf", "@type": "@id"},
	"cutFrom": {"@id": "tribo:cutFrom", "@type": "@id"},
	"performedOn": {"@id": "tribo:performedOn", "@type": "@id"},
	"performedBy": {"@id": "tribo:performedBy", "@type": "@id"},
	"duringRun": {"@id": "tribo:duringRun", "@type": "@id"},
	"started_at": {"@id": "prov:startedAtTime"},
}

MEASUREMENTS = [
	# (csv file, json key, @type, instrument id, points key)
	("xrf_map.csv", "xrf", "CompositionMeasurement", "m4_tornado", "points"),
	("profilometry.csv", "profilometry", "ProfilometryMeasurement", "profilometer", "points"),
	("nanoindentation.csv", "nanoindentation", "NanoindentationMeasurement", "ti980", "indents"),
]


class JsonConverter:
	def __init__(self, cfg: dict, csv_root: Path, json_root: Path) -> None:
		self.cfg = cfg
		self.json_root = json_root
		self.reader = CorpusReader(csv_root)
		self.manifest = self.reader.manifest
		self.graph: list[dict] = []  # hybrid metadata entities
		self.full_bytes = 0
		self.full_files = 0
		self.counts = {"coupons": 0, "tracks": 0, "cycle_rows": 0, "loop_rows": 0, "xrf_rows": 0}

	def rows_as_dicts(self, relpath: str) -> list[dict]:
		return self.reader.dicts(relpath)

	def source_ref(self, relpath: str) -> dict:
		rows, _, sha = self.manifest[relpath]
		return {"sourceFile": relpath, "rows": rows, "sha256": sha}

	# --- flat entities (hybrid graph) ---

	def convert_flat(self) -> None:
		for inst in self.cfg["instruments"]:
			self.graph.append({
				"@id": f"instrument:{inst['instrument_id']}",
				"@type": "Instrument",
				"name": inst["name"],
				"role": inst["role"],
			})
		for row in self.rows_as_dicts("batches.csv"):
			self.graph.append({"@id": f"batch:{row['batch_code']}", "@type": "Batch", **row})
			self.graph.append({
				"@id": f"batch:{row['batch_code']}-simtra",
				"@type": "SimtraProfile",
				"performedOn": f"batch:{row['batch_code']}",
				"performedBy": "instrument:simtra",
				**self.source_ref(f"simtra/simtra_profile_{row['batch_code']}.csv"),
			})
		for row in self.rows_as_dicts("runs.csv"):
			self.graph.append({"@id": f"run:{row['run_code']}", "@type": "Run", **row})

	# --- per-coupon conversion ---

	def convert_coupons(self) -> None:
		h = self.cfg["hierarchy"]
		full_dir = self.json_root / "full"
		full_dir.mkdir(parents=True, exist_ok=True)
		for batch in self.cfg["deposition_matrix"]:
			b_code = batch["batch_code"]
			for w in range(1, h["wafers_per_batch"] + 1):
				wafer_code = f"{b_code}-W{w}"
				wafer_dir = f"batch_{b_code}/wafer_W{w}"
				info = self.rows_as_dicts(f"{wafer_dir}/wafer_info.csv")[0]
				self.graph.append({
					"@id": f"wafer:{wafer_code}", "@type": "Wafer",
					"partOf": f"batch:{b_code}", **info,
				})
				for c in range(1, h["coupons_per_wafer"] + 1):
					self.convert_coupon(f"{wafer_dir}/coupon_C{c:02d}", full_dir)
				logger.debug("wafer %s converted", wafer_code)
			logger.info("batch %s converted: %.1f MiB full so far", b_code, self.full_bytes / (1024 * 1024))

	def convert_coupon(self, rel_dir: str, full_dir: Path) -> None:
		info = self.rows_as_dicts(f"{rel_dir}/coupon_info.csv")[0]
		code = info["coupon_code"]
		coupon_id = f"coupon:{code}"
		is_friction = info["run_code"] != "RESERVE"

		full_node: dict = {
			"@context": CONTEXT, "@id": coupon_id, "@type": "Coupon",
			"cutFrom": f"wafer:{info['wafer_code']}", **info,
		}
		hybrid_node: dict = {
			"@id": coupon_id, "@type": "Coupon",
			"cutFrom": f"wafer:{info['wafer_code']}", **info,
		}

		for csv_name, key, cls, inst, points_key in MEASUREMENTS:
			relpath = f"{rel_dir}/{csv_name}"
			meta = {"@type": cls, "performedOn": coupon_id, "performedBy": f"instrument:{inst}"}
			full_node[key] = {**meta, points_key: self.rows_as_dicts(relpath)}
			hybrid_node[key] = {**meta, **self.source_ref(relpath)}
		self.counts["xrf_rows"] += self.manifest[f"{rel_dir}/xrf_map.csv"][0]

		afm = self.rows_as_dicts(f"{rel_dir}/afm.csv")[0]
		afm_node = {"@type": "AFMMeasurement", "performedOn": coupon_id, "performedBy": "instrument:afm", **afm}
		full_node["afm"] = afm_node
		hybrid_node["afm"] = afm_node

		if is_friction:
			tracks_full = []
			for t in range(1, self.cfg["friction_assignment"]["tracks_per_friction_coupon"] + 1):
				full_track, hybrid_track = self.convert_track(f"{rel_dir}/track_T{t}", coupon_id, info["run_code"])
				tracks_full.append(full_track)
				self.graph.append(hybrid_track)
			full_node["tracks"] = tracks_full

		out_path = full_dir / f"{code}.jsonld"
		text = json.dumps(full_node, separators=(",", ":"))
		data = text.encode("ascii")
		out_path.write_bytes(data)
		self.full_bytes += len(data)
		self.full_files += 1
		self.graph.append(hybrid_node)
		self.counts["coupons"] += 1

	def convert_track(self, rel_dir: str, coupon_id: str, run_code: str) -> tuple[dict, dict]:
		info = self.rows_as_dicts(f"{rel_dir}/track_info.csv")[0]
		track_id = f"track:{info['track_code']}"
		base = {
			"@id": track_id, "@type": "FrictionTest",
			"partOf": coupon_id, "duringRun": f"run:{run_code}",
			"performedBy": "instrument:rapid", **info,
		}
		wear = self.rows_as_dicts(f"{rel_dir}/wear.csv")[0]
		wear_node = {"@type": "WearMeasurement", **wear}

		cycles_path = f"{rel_dir}/cof_vs_cycle.csv"
		loops_path = f"{rel_dir}/friction_loops.csv"
		full_track = {
			**base,
			"cycles": self.rows_as_dicts(cycles_path),
			"loop_points": self.rows_as_dicts(loops_path),
			"wear": wear_node,
		}
		hybrid_track = {
			**base,
			"cycles": self.source_ref(cycles_path),
			"loop_points": self.source_ref(loops_path),
			"wear": wear_node,
		}
		self.counts["cycle_rows"] += len(full_track["cycles"])
		self.counts["loop_rows"] += len(full_track["loop_points"])
		self.counts["tracks"] += 1
		return full_track, hybrid_track

	# --- outputs ---

	def write_hybrid(self) -> int:
		hybrid_dir = self.json_root / "hybrid"
		hybrid_dir.mkdir(parents=True, exist_ok=True)
		path = hybrid_dir / "dataset.jsonld"
		with open(path, "w", encoding="ascii", newline="\n") as fh:
			json.dump({"@context": CONTEXT, "@graph": self.graph}, fh, separators=(",", ":"))
		return path.stat().st_size

	def validate_counts(self) -> None:
		v = self.cfg["volumes"]
		expected = {
			"coupons": v["coupons_total"],
			"tracks": v["tracks_total"],
			"cycle_rows": v["cycle_rows_total"],
			"loop_rows": v["loop_points_total"],
			"xrf_rows": v["xrf_points_total"],
		}
		for key, exp in expected.items():
			if self.counts[key] != exp:
				raise ValidationError(f"count mismatch: {key}: converted {self.counts[key]} != expected {exp}")
		logger.info("all converted counts match lab_config volumes: %s", self.counts)


def validate_sample(paths: list[Path]) -> None:
	"""Smoke test (spec 04 + test-pipeline-validation): json.load + rdflib parse."""
	for path in paths:
		with open(path, encoding="ascii") as fh:
			json.load(fh)
		g = Graph()
		g.parse(path, format="json-ld")
		if len(g) == 0:
			raise ValidationError(f"{path}: rdflib parsed 0 triples")
		logger.info("sample %s: valid JSON, %d triples via rdflib", path.name, len(g))


def main() -> int:
	args = parse_task_args("Task 04: convert the CSV corpus to JSON-LD")
	out_root: Path = args.out_root
	json_root = out_root / "json"
	try:
		check_dependencies(out_root, DEPENDS_ON)
		remove_stale_marker(out_root, TASK_ID)
		if json_root.exists():
			logger.info("re-run: removing previous output %s", json_root)
			shutil.rmtree(json_root)
		cfg = load_lab_config(out_root)

		conv = JsonConverter(cfg, out_root / "csv", json_root)
		conv.convert_flat()
		conv.convert_coupons()
		conv.validate_counts()
		hybrid_bytes = conv.write_hybrid()
		validate_sample([
			json_root / "full" / "B722-W2-C13.jsonld",
			json_root / "hybrid" / "dataset.jsonld",
		])
		update_storage_sizes(out_root / "bench", "json", [("full", conv.full_bytes), ("hybrid", hybrid_bytes)])
		entries = {
			"full_files": conv.full_files,
			"full_bytes": conv.full_bytes,
			"hybrid_bytes": hybrid_bytes,
			"coupons": conv.counts["coupons"],
			"tracks": conv.counts["tracks"],
			"cycle_rows": conv.counts["cycle_rows"],
			"loop_points": conv.counts["loop_rows"],
		}
		entries.update(process_metrics())
		marker_path = write_marker(out_root, TASK_ID, entries)
	except Exception:
		logger.critical("task %s failed", TASK_ID, exc_info=True)
		return 1

	print(f"task 04 ok: full {conv.full_files} files, {conv.full_bytes / (1024 * 1024):.1f} MiB; "
		f"hybrid {hybrid_bytes / (1024 * 1024):.1f} MiB")
	print(f"entities: coupons={conv.counts['coupons']} tracks={conv.counts['tracks']} "
		f"cycles={conv.counts['cycle_rows']} loop_points={conv.counts['loop_rows']}")
	print(f"marker: {marker_path}")
	return 0


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