"""Task 07 - Convert the CSV corpus to RDF.

Streams the corpus into ./out/rdf/dataset.nt.gz (N-Triples, gzip) and
./out/rdf/dataset.ttl (Turtle) with the task 04 vocabulary (tribo# + PROV-O
for started_at), then bulk-loads the N-Triples into an embedded on-disk
oxigraph store (./out/rdf/oxigraph_store/) and smoke-tests it with SPARQL
COUNT queries. Bulk points (cycles, loop points, map points) carry no
rdf:type - the containing predicate types them (documented decision per
spec 07), which is why the triple count lands below the spec's 25-45M
sketch. Appends serialization and store sizes to storage_sizes.csv and
writes ./out/.done/07.ok. Spec: docs/specs/07_convert_rdf.md.
"""

from __future__ import annotations

import gzip
import io
import logging
import shutil
import sys
import time
from pathlib import Path

from pyoxigraph import Store

try:
	from pyoxigraph import RdfFormat
	NT_FORMAT = RdfFormat.N_TRIPLES
except ImportError:  # pyoxigraph < 0.5 uses MIME strings
	NT_FORMAT = "application/n-triples"

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 = "07"
DEPENDS_ON = ["01", "03"]

VOCAB = "https://sandia.gov/ontology/tribology#"
ID_BASE = "https://sandia.gov/ontology/tribology/id/"
XSD = "http://www.w3.org/2001/XMLSchema#"

PREFIXES = {
	"tribo": VOCAB,
	"batch": ID_BASE + "batch/",
	"wafer": ID_BASE + "wafer/",
	"coupon": ID_BASE + "coupon/",
	"track": ID_BASE + "track/",
	"run": ID_BASE + "run/",
	"instrument": ID_BASE + "instrument/",
	"simtra": ID_BASE + "simtra/",
	"prov": "http://www.w3.org/ns/prov#",
	"xsd": XSD,
}

INT_COLS = {
	"cycle", "grid_x", "grid_y", "indent_id", "plates", "plate", "probe", "square",
	"row_no", "wafer_index", "coupons", "friction_coupons", "reserve_coupons",
	"pt_gun_tilt_deg", "au_gun_tilt_deg", "pt_power_w", "au_power_w",
	"pt_discharge_v", "au_discharge_v",
}
DOUBLE_COLS = {
	"cof", "position_um", "friction_force_mn", "pt_wtpct", "au_wtpct", "thickness_um",
	"x_um", "y_um", "hardness_gpa", "reduced_modulus_gpa", "max_load_mn", "ra_nm",
	"rq_nm", "wear_volume_um3", "k_archard", "sliding_distance_m", "load_mn",
	"stroke_mm", "speed_mm_s", "angle_deg", "energy_ev", "pt_flux", "au_flux",
	"au_wtpct_mean",
}
DATE_COLS = {"date", "deposition_date"}

MEASUREMENTS = [
	("xrf_map.csv", "xrf", "CompositionMeasurement", "m4_tornado"),
	("profilometry.csv", "profilometry", "ProfilometryMeasurement", "profilometer"),
	("nanoindentation.csv", "nanoindentation", "NanoindentationMeasurement", "ti980"),
]


def lit(col: str, token: str) -> tuple:
	if col in INT_COLS:
		return ("l", token, "integer")
	if col in DOUBLE_COLS:
		return ("l", token, "double")
	if col in DATE_COLS:
		return ("l", token, "date")
	if col == "started_at":
		return ("l", token, "dateTime")
	return ("l", token, None)


def iri(curie: str) -> tuple:
	return ("i", curie)


def node(pairs: list) -> tuple:
	return ("n", pairs)


class RdfEmitter:
	"""Writes the same subject blocks to N-Triples and Turtle streams."""

	def __init__(self, nt_fh, ttl_fh) -> None:
		self.nt = nt_fh
		self.ttl = ttl_fh
		self.triples = 0
		self._blank = 0
		for prefix, base in PREFIXES.items():
			self.ttl.write(f"@prefix {prefix}: <{base}> .\n")
		self.ttl.write("\n")

	def _full(self, curie: str) -> str:
		prefix, local = curie.split(":", 1)
		return f"<{PREFIXES[prefix]}{local}>"

	def _pred_nt(self, pred: str) -> str:
		if pred == "a":
			return "<http://www.w3.org/1999/02/22-rdf-syntax-ns#type>"
		return self._full(pred)

	@staticmethod
	def _escape(text: str) -> str:
		return text.replace("\\", "\\\\").replace('"', '\\"')

	def _obj_nt(self, obj: tuple) -> str:
		kind = obj[0]
		if kind == "i":
			return self._full(obj[1])
		if kind == "l":
			_, token, dtype = obj
			if dtype is None:
				return f'"{self._escape(token)}"'
			return f'"{token}"^^<{XSD}{dtype}>'
		self._blank += 1
		label = f"_:b{self._blank}"
		for pred, sub_obj in obj[1]:
			self.nt.write(f"{label} {self._pred_nt(pred)} {self._obj_nt(sub_obj)} .\n")
			self.triples += 1
		return label

	def _obj_ttl(self, obj: tuple) -> str:
		kind = obj[0]
		if kind == "i":
			return obj[1]
		if kind == "l":
			_, token, dtype = obj
			if dtype is None:
				return f'"{self._escape(token)}"'
			if dtype == "integer":
				return token
			return f'"{token}"^^xsd:{dtype}'
		inner = " ; ".join(f"{p} {self._obj_ttl(o)}" for p, o in obj[1])
		return f"[ {inner} ]"

	def emit(self, subject: str, pairs: list) -> None:
		subj_nt = self._full(subject)
		ttl_parts = []
		for pred, obj in pairs:
			# N-Triples: nested nodes expand to labelled blanks first, so the
			# referencing triple is written after the node's own triples.
			obj_nt = self._obj_nt(obj)
			self.nt.write(f"{subj_nt} {self._pred_nt(pred)} {obj_nt} .\n")
			self.triples += 1
			ttl_parts.append(f"{pred} {self._obj_ttl(obj)}")
		self.ttl.write(subject + " " + " ;\n  ".join(ttl_parts) + " .\n")


class RdfConverter:
	def __init__(self, cfg: dict, csv_root: Path, emitter: RdfEmitter) -> None:
		self.cfg = cfg
		self.reader = CorpusReader(csv_root)
		self.em = emitter
		self.counts = {"coupons": 0, "tracks": 0, "cycle_nodes": 0, "loop_nodes": 0, "xrf_nodes": 0}

	def raw_rows(self, relpath: str) -> list[dict[str, str]]:
		header, rows = self.reader.read(relpath)
		keys = [h.lower() for h in header]
		return [dict(zip(keys, row)) for row in rows]

	def prop_pairs(self, row: dict[str, str], skip: tuple = ()) -> list:
		return [
			("prov:startedAtTime" if col == "started_at" else f"tribo:{col}", lit(col, token))
			for col, token in row.items()
			if col not in skip and token != ""
		]

	def convert(self) -> None:
		for inst in self.cfg["instruments"]:
			self.em.emit(f"instrument:{inst['instrument_id']}", [
				("a", iri("tribo:Instrument")),
				("tribo:name", ("l", inst["name"], None)),
				("tribo:role", ("l", inst["role"], None)),
			])
		for row in self.raw_rows("batches.csv"):
			code = row["batch_code"]
			self.em.emit(f"batch:{code}", [("a", iri("tribo:Batch"))] + self.prop_pairs(row))
			points = [
				("tribo:hasPoint", node([("tribo:row_no", ("l", str(no), "integer"))] + self.prop_pairs(p)))
				for no, p in enumerate(self.raw_rows(f"simtra/simtra_profile_{code}.csv"), 1)
			]
			self.em.emit(f"simtra:{code}", [
				("a", iri("tribo:SimtraProfile")),
				("tribo:performedOn", iri(f"batch:{code}")),
				("tribo:performedBy", iri("instrument:simtra")),
			] + points)
		for row in self.raw_rows("runs.csv"):
			self.em.emit(f"run:{row['run_code']}", [("a", iri("tribo:Run"))] + self.prop_pairs(row))

		h = self.cfg["hierarchy"]
		for batch in self.cfg["deposition_matrix"]:
			b_code = batch["batch_code"]
			for w in range(1, h["wafers_per_batch"] + 1):
				wafer_dir = f"batch_{b_code}/wafer_W{w}"
				info = self.raw_rows(f"{wafer_dir}/wafer_info.csv")[0]
				self.em.emit(f"wafer:{info['wafer_code']}", [
					("a", iri("tribo:Wafer")),
					("tribo:partOf", iri(f"batch:{b_code}")),
				] + self.prop_pairs(info, skip=("batch_code",)))
				for c in range(1, h["coupons_per_wafer"] + 1):
					self.convert_coupon(f"{wafer_dir}/coupon_C{c:02d}")
			logger.info("batch %s converted (%d triples so far)", b_code, self.em.triples)

	def convert_coupon(self, rel_dir: str) -> None:
		info = self.raw_rows(f"{rel_dir}/coupon_info.csv")[0]
		code = info["coupon_code"]
		subject = f"coupon:{code}"
		is_friction = info["run_code"] != "RESERVE"
		pairs = [("a", iri("tribo:Coupon")), ("tribo:cutFrom", iri(f"wafer:{info['wafer_code']}"))]
		pairs += self.prop_pairs(info, skip=("batch_code", "wafer_code", "run_code"))
		pairs.append(("tribo:run_code", ("l", info["run_code"], None)))
		if is_friction:
			pairs.append(("tribo:duringRun", iri(f"run:{info['run_code']}")))

		for csv_name, key, cls, inst in MEASUREMENTS:
			points = [
				("tribo:hasPoint" if key != "nanoindentation" else "tribo:hasIndent", node(self.prop_pairs(p)))
				for p in self.raw_rows(f"{rel_dir}/{csv_name}")
			]
			if key == "xrf":
				self.counts["xrf_nodes"] += len(points)
			pairs.append((f"tribo:{key}", node([
				("a", iri(f"tribo:{cls}")),
				("tribo:performedOn", iri(subject)),
				("tribo:performedBy", iri(f"instrument:{inst}")),
			] + points)))
		afm = self.raw_rows(f"{rel_dir}/afm.csv")[0]
		pairs.append(("tribo:afm", node([
			("a", iri("tribo:AFMMeasurement")),
			("tribo:performedOn", iri(subject)),
			("tribo:performedBy", iri("instrument:afm")),
		] + self.prop_pairs(afm))))
		self.em.emit(subject, pairs)
		self.counts["coupons"] += 1

		if is_friction:
			for t in range(1, self.cfg["friction_assignment"]["tracks_per_friction_coupon"] + 1):
				self.convert_track(f"{rel_dir}/track_T{t}", subject)

	def convert_track(self, rel_dir: str, coupon_subject: str) -> None:
		info = self.raw_rows(f"{rel_dir}/track_info.csv")[0]
		pairs = [
			("a", iri("tribo:FrictionTest")),
			("tribo:partOf", iri(coupon_subject)),
			("tribo:duringRun", iri(f"run:{info['run_code']}")),
			("tribo:performedBy", iri("instrument:rapid")),
		] + self.prop_pairs(info, skip=("track_code", "run_code"))
		pairs.insert(1, ("tribo:track_code", ("l", info["track_code"], None)))

		cycles = self.raw_rows(f"{rel_dir}/cof_vs_cycle.csv")
		pairs += [("tribo:hasCycle", node(self.prop_pairs(row))) for row in cycles]
		self.counts["cycle_nodes"] += len(cycles)
		loops = self.raw_rows(f"{rel_dir}/friction_loops.csv")
		pairs += [("tribo:hasLoopPoint", node(self.prop_pairs(row))) for row in loops]
		self.counts["loop_nodes"] += len(loops)
		wear = self.raw_rows(f"{rel_dir}/wear.csv")[0]
		pairs.append(("tribo:wear", node([("a", iri("tribo:WearMeasurement"))] + self.prop_pairs(wear))))
		self.em.emit(f"track:{info['track_code']}", pairs)
		self.counts["tracks"] += 1

	def validate_counts(self) -> None:
		v = self.cfg["volumes"]
		expected = {
			"coupons": v["coupons_total"],
			"tracks": v["tracks_total"],
			"cycle_nodes": v["cycle_rows_total"],
			"loop_nodes": v["loop_points_total"],
			"xrf_nodes": 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 bulk_load(store: Store, nt_gz: Path) -> None:
	with gzip.open(nt_gz, "rb") as fh:
		try:
			store.bulk_load(fh, NT_FORMAT)
		except TypeError:
			store.bulk_load(input=fh, format=NT_FORMAT)


def sparql_count(store: Store, query: str) -> int:
	solutions = store.query(query)
	return int(next(iter(solutions))[0].value)


def main() -> int:
	args = parse_task_args("Task 07: convert the CSV corpus to RDF")
	out_root: Path = args.out_root
	rdf_root = out_root / "rdf"
	nt_gz = rdf_root / "dataset.nt.gz"
	ttl_path = rdf_root / "dataset.ttl"
	store_dir = rdf_root / "oxigraph_store"
	try:
		check_dependencies(out_root, DEPENDS_ON)
		remove_stale_marker(out_root, TASK_ID)
		if rdf_root.exists():
			logger.info("re-run: removing previous output %s", rdf_root)
			shutil.rmtree(rdf_root)
		rdf_root.mkdir(parents=True)
		cfg = load_lab_config(out_root)

		# mtime=0 keeps the gzip byte-identical across re-runs (data-determinism)
		with open(nt_gz, "wb") as raw:
			gz = gzip.GzipFile(filename="", mode="wb", fileobj=raw, compresslevel=6, mtime=0)
			with io.TextIOWrapper(gz, encoding="ascii", newline="\n") as nt_fh, \
					open(ttl_path, "w", encoding="ascii", newline="\n") as ttl_fh:
				emitter = RdfEmitter(nt_fh, ttl_fh)
				conv = RdfConverter(cfg, out_root / "csv", emitter)
				conv.convert()
				conv.validate_counts()
				triples = emitter.triples
		nt_bytes = nt_gz.stat().st_size
		ttl_bytes = ttl_path.stat().st_size
		logger.info("serialized %d triples: nt.gz %.1f MiB, ttl %.1f MiB",
			triples, nt_bytes / 1048576, ttl_bytes / 1048576)

		t0 = time.time()
		store = Store(str(store_dir))
		bulk_load(store, nt_gz)
		store.flush()
		load_seconds = round(time.time() - t0, 1)

		total = sparql_count(store, "SELECT (COUNT(*) AS ?n) WHERE { ?s ?p ?o }")
		if total != triples:
			raise ValidationError(f"store holds {total} triples, serialized {triples}")
		q1 = sparql_count(
			store,
			f"SELECT (COUNT(?c) AS ?n) WHERE {{ <{ID_BASE}track/B722-W2-C13-T2> <{VOCAB}hasCycle> ?c }}",
		)
		if q1 != cfg["volumes"]["cycles_per_track"]:
			raise ValidationError(f"SPARQL Q1 smoke test returned {q1} cycles")
		logger.info("store smoke tests ok: %d triples, Q1 track has %d cycles (load %.1f s)", total, q1, load_seconds)
		del store
		store_bytes = sum(f.stat().st_size for f in store_dir.rglob("*") if f.is_file())

		update_storage_sizes(out_root / "bench", "rdf", [
			("nt_gz", nt_bytes), ("ttl", ttl_bytes), ("store", store_bytes),
		])
		entries = {
			"triples": triples,
			"nt_gz_bytes": nt_bytes,
			"ttl_bytes": ttl_bytes,
			"store_bytes": store_bytes,
			"store_load_seconds": load_seconds,
			"coupons": conv.counts["coupons"],
			"tracks": conv.counts["tracks"],
			"cycle_nodes": conv.counts["cycle_nodes"],
			"loop_nodes": conv.counts["loop_nodes"],
		}
		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 07 ok: {triples} triples; nt.gz {nt_bytes / 1048576:.1f} MiB, "
		f"ttl {ttl_bytes / 1048576:.1f} MiB, store {store_bytes / 1048576:.1f} MiB "
		f"(bulk load {load_seconds} s)")
	print(f"entities: coupons={conv.counts['coupons']} tracks={conv.counts['tracks']} "
		f"cycles={conv.counts['cycle_nodes']} loop_points={conv.counts['loop_nodes']}")
	print(f"marker: {marker_path}")
	return 0


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