"""Task 05 - Convert the CSV corpus to SQLite.

Builds ./out/sqlite/tribo.db: variant B key schema (composite natural PKs
on bulk tables, WITHOUT ROWID, enforced foreign keys - see
docs/research/Tribology_FK_Architecture.html), Q1-Q7 indexes created after
the load, track_summary precomputed. Every consumed CSV is verified
against MANIFEST.csv. Appends the db size to storage_sizes.csv and writes
./out/.done/05.ok. Spec: docs/specs/05_convert_sqlite.md.
"""

from __future__ import annotations

import logging
import sqlite3
import sys
from pathlib import Path

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

logger = logging.getLogger(__name__)

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

BATCH_ROWS = 50000  # rows per executemany transaction, spec 05
LOAD_PRAGMAS = [
	"PRAGMA journal_mode=OFF",
	"PRAGMA synchronous=OFF",
	"PRAGMA cache_size=-262144",
	"PRAGMA foreign_keys=ON",
]

DDL = [
	"""CREATE TABLE instruments (
		instrument_id INTEGER PRIMARY KEY,
		instrument_code TEXT NOT NULL UNIQUE,
		name TEXT NOT NULL,
		role TEXT NOT NULL
	)""",
	"""CREATE TABLE batches (
		batch_id INTEGER PRIMARY KEY,
		batch_code TEXT NOT NULL UNIQUE,
		pt_gun_tilt_deg INTEGER NOT NULL,
		au_gun_tilt_deg INTEGER NOT NULL,
		pt_power_w INTEGER NOT NULL,
		au_power_w INTEGER NOT NULL,
		pt_discharge_v INTEGER NOT NULL,
		au_discharge_v INTEGER NOT NULL,
		deposition_date TEXT NOT NULL
	)""",
	"""CREATE TABLE runs (
		run_id INTEGER PRIMARY KEY,
		run_code TEXT NOT NULL UNIQUE,
		environment TEXT NOT NULL,
		date TEXT NOT NULL,
		plates INTEGER NOT NULL,
		operator TEXT NOT NULL
	)""",
	"""CREATE TABLE wafers (
		wafer_id INTEGER PRIMARY KEY,
		wafer_code TEXT NOT NULL UNIQUE,
		batch_id INTEGER NOT NULL REFERENCES batches(batch_id),
		wafer_index INTEGER NOT NULL,
		deposition_date TEXT NOT NULL,
		coupons INTEGER NOT NULL,
		friction_coupons INTEGER NOT NULL,
		reserve_coupons INTEGER NOT NULL
	)""",
	"""CREATE TABLE coupons (
		coupon_id INTEGER PRIMARY KEY,
		coupon_code TEXT NOT NULL UNIQUE,
		wafer_id INTEGER NOT NULL REFERENCES wafers(wafer_id),
		batch_id INTEGER NOT NULL REFERENCES batches(batch_id),
		grid_row INTEGER NOT NULL,
		grid_col INTEGER NOT NULL,
		thickness_um REAL NOT NULL,
		ra_nm REAL NOT NULL,
		au_wtpct_mean REAL NOT NULL,
		run_id INTEGER REFERENCES runs(run_id),
		plate INTEGER,
		probe INTEGER,
		square INTEGER
	)""",
	"""CREATE TABLE tracks (
		track_id INTEGER PRIMARY KEY,
		track_code TEXT NOT NULL UNIQUE,
		coupon_id INTEGER NOT NULL REFERENCES coupons(coupon_id),
		run_id INTEGER NOT NULL REFERENCES runs(run_id),
		environment TEXT NOT NULL,
		load_mn REAL NOT NULL,
		stroke_mm REAL NOT NULL,
		speed_mm_s REAL NOT NULL,
		counterface_id TEXT NOT NULL,
		started_at TEXT NOT NULL
	)""",
	"""CREATE TABLE simtra_profiles (
		batch_id INTEGER NOT NULL REFERENCES batches(batch_id),
		row_no INTEGER NOT NULL,
		angle_deg REAL NOT NULL,
		energy_ev REAL NOT NULL,
		pt_flux REAL NOT NULL,
		au_flux REAL NOT NULL,
		PRIMARY KEY (batch_id, row_no)
	) WITHOUT ROWID""",
	"""CREATE TABLE xrf_points (
		coupon_id INTEGER NOT NULL REFERENCES coupons(coupon_id),
		grid_x INTEGER NOT NULL,
		grid_y INTEGER NOT NULL,
		pt_wtpct REAL NOT NULL,
		au_wtpct REAL NOT NULL,
		PRIMARY KEY (coupon_id, grid_x, grid_y)
	) WITHOUT ROWID""",
	"""CREATE TABLE profilometry_points (
		coupon_id INTEGER NOT NULL REFERENCES coupons(coupon_id),
		grid_x INTEGER NOT NULL,
		grid_y INTEGER NOT NULL,
		thickness_um REAL NOT NULL,
		PRIMARY KEY (coupon_id, grid_x, grid_y)
	) WITHOUT ROWID""",
	"""CREATE TABLE nanoindentation (
		coupon_id INTEGER NOT NULL REFERENCES coupons(coupon_id),
		indent_id INTEGER NOT NULL,
		x_um REAL NOT NULL,
		y_um REAL NOT NULL,
		hardness_gpa REAL NOT NULL,
		reduced_modulus_gpa REAL NOT NULL,
		max_load_mn REAL NOT NULL,
		PRIMARY KEY (coupon_id, indent_id)
	) WITHOUT ROWID""",
	"""CREATE TABLE afm (
		coupon_id INTEGER PRIMARY KEY REFERENCES coupons(coupon_id),
		ra_nm REAL NOT NULL,
		rq_nm REAL NOT NULL,
		image_file TEXT NOT NULL
	)""",
	"""CREATE TABLE friction_cycles (
		track_id INTEGER NOT NULL REFERENCES tracks(track_id),
		cycle INTEGER NOT NULL,
		cof REAL NOT NULL,
		PRIMARY KEY (track_id, cycle)
	) WITHOUT ROWID""",
	"""CREATE TABLE friction_loop_points (
		track_id INTEGER NOT NULL REFERENCES tracks(track_id),
		cycle INTEGER NOT NULL,
		pt INTEGER NOT NULL,
		position_um REAL NOT NULL,
		friction_force_mn REAL NOT NULL,
		PRIMARY KEY (track_id, cycle, pt)
	) WITHOUT ROWID""",
	"""CREATE TABLE wear (
		track_id INTEGER PRIMARY KEY REFERENCES tracks(track_id),
		wear_volume_um3 REAL NOT NULL,
		k_archard REAL NOT NULL,
		sliding_distance_m REAL NOT NULL
	)""",
	"""CREATE TABLE track_summary (
		track_id INTEGER PRIMARY KEY REFERENCES tracks(track_id),
		cof_ss_mean REAL NOT NULL,
		cof_ss_std REAL NOT NULL,
		run_in_cycles INTEGER NOT NULL
	)""",
]

INDEX_DDL = [
	"CREATE INDEX ix_coupons_au_wtpct_mean ON coupons(au_wtpct_mean)",  # Q2 selective filter
	"CREATE INDEX ix_tracks_run_id ON tracks(run_id)",  # Q5/Q7 joins to runs
	"CREATE INDEX ix_tracks_coupon_id ON tracks(coupon_id)",  # Q3/Q6 joins to coupons
	"CREATE INDEX ix_runs_environment ON runs(environment)",  # Q7 grouping
	"CREATE INDEX ix_tracks_environment_load ON tracks(environment, load_mn)",  # Q4 filter
]

INSERTS = {
	"instruments": "INSERT INTO instruments VALUES (?,?,?,?)",
	"batches": "INSERT INTO batches VALUES (?,?,?,?,?,?,?,?,?)",
	"runs": "INSERT INTO runs VALUES (?,?,?,?,?,?)",
	"wafers": "INSERT INTO wafers VALUES (?,?,?,?,?,?,?,?)",
	"coupons": "INSERT INTO coupons VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?)",
	"tracks": "INSERT INTO tracks VALUES (?,?,?,?,?,?,?,?,?,?)",
	"simtra_profiles": "INSERT INTO simtra_profiles VALUES (?,?,?,?,?,?)",
	"xrf_points": "INSERT INTO xrf_points VALUES (?,?,?,?,?)",
	"profilometry_points": "INSERT INTO profilometry_points VALUES (?,?,?,?)",
	"nanoindentation": "INSERT INTO nanoindentation VALUES (?,?,?,?,?,?,?)",
	"afm": "INSERT INTO afm VALUES (?,?,?,?)",
	"friction_cycles": "INSERT INTO friction_cycles VALUES (?,?,?)",
	"friction_loop_points": "INSERT INTO friction_loop_points VALUES (?,?,?,?,?)",
	"wear": "INSERT INTO wear VALUES (?,?,?,?)",
	"track_summary": "INSERT INTO track_summary VALUES (?,?,?,?)",
}


class SqliteLoader:
	def __init__(self, cfg: dict, csv_root: Path, db_path: Path) -> None:
		self.cfg = cfg
		self.reader = CorpusReader(csv_root)
		db_path.parent.mkdir(parents=True, exist_ok=True)
		self.conn = sqlite3.connect(db_path)
		for pragma in LOAD_PRAGMAS:
			self.conn.execute(pragma)
		for ddl in DDL:
			self.conn.execute(ddl)
		self.buffers: dict[str, list[tuple]] = {}
		self.batch_ids: dict[str, int] = {}
		self.run_ids: dict[str, int] = {}

	def put(self, table: str, row: tuple) -> None:
		buf = self.buffers.setdefault(table, [])
		buf.append(row)
		if len(buf) >= BATCH_ROWS:
			self.flush(table)

	def flush(self, table: str) -> None:
		buf = self.buffers.get(table)
		if buf:
			self.conn.executemany(INSERTS[table], buf)
			self.conn.commit()
			buf.clear()

	def flush_all(self) -> None:
		for table in list(self.buffers):
			self.flush(table)

	# --- dimension + bulk loading, in FK dependency order ---

	def load_flat(self) -> None:
		for i, inst in enumerate(self.cfg["instruments"], 1):
			self.put("instruments", (i, inst["instrument_id"], inst["name"], inst["role"]))
		for i, row in enumerate(self.reader.dicts("batches.csv"), 1):
			self.batch_ids[row["batch_code"]] = i
			self.put("batches", (
				i, row["batch_code"], row["pt_gun_tilt_deg"], row["au_gun_tilt_deg"],
				row["pt_power_w"], row["au_power_w"], row["pt_discharge_v"],
				row["au_discharge_v"], row["deposition_date"],
			))
		for i, row in enumerate(self.reader.dicts("runs.csv"), 1):
			self.run_ids[row["run_code"]] = i
			self.put("runs", (i, row["run_code"], row["environment"], row["date"], row["plates"], row["operator"]))
		self.flush_all()
		for code, batch_id in self.batch_ids.items():
			for row_no, row in enumerate(self.reader.dicts(f"simtra/simtra_profile_{code}.csv"), 1):
				self.put("simtra_profiles", (batch_id, row_no, row["angle_deg"], row["energy_ev"], row["pt_flux"], row["au_flux"]))

	def load_hierarchy(self) -> None:
		h = self.cfg["hierarchy"]
		fa = self.cfg["friction_assignment"]
		wafer_id = 0
		coupon_id = 0
		track_id = 0
		for batch in self.cfg["deposition_matrix"]:
			b_code = batch["batch_code"]
			batch_id = self.batch_ids[b_code]
			for w in range(1, h["wafers_per_batch"] + 1):
				wafer_id += 1
				wafer_dir = f"batch_{b_code}/wafer_W{w}"
				info = self.reader.dicts(f"{wafer_dir}/wafer_info.csv")[0]
				self.put("wafers", (
					wafer_id, info["wafer_code"], batch_id, info["wafer_index"],
					info["deposition_date"], info["coupons"], info["friction_coupons"], info["reserve_coupons"],
				))
				self.flush("wafers")
				for c in range(1, h["coupons_per_wafer"] + 1):
					coupon_id += 1
					rel_dir = f"{wafer_dir}/coupon_C{c:02d}"
					track_id = self.load_coupon(rel_dir, coupon_id, wafer_id, batch_id, track_id, fa)
			logger.info("batch %s loaded (through coupon %d, track %d)", b_code, coupon_id, track_id)
		self.flush_all()

	def load_coupon(self, rel_dir: str, coupon_id: int, wafer_id: int, batch_id: int, track_id: int, fa: dict) -> int:
		info = self.reader.dicts(f"{rel_dir}/coupon_info.csv")[0]
		is_friction = info["run_code"] != "RESERVE"
		run_id = self.run_ids[info["run_code"]] if is_friction else None
		self.put("coupons", (
			coupon_id, info["coupon_code"], wafer_id, batch_id,
			info["grid_row"], info["grid_col"], info["thickness_um"], info["ra_nm"],
			info["au_wtpct_mean"], run_id,
			info.get("plate"), info.get("probe"), info.get("square"),
		))
		self.flush("coupons")

		for row in self.reader.dicts(f"{rel_dir}/xrf_map.csv"):
			self.put("xrf_points", (coupon_id, row["grid_x"], row["grid_y"], row["pt_wtpct"], row["au_wtpct"]))
		for row in self.reader.dicts(f"{rel_dir}/profilometry.csv"):
			self.put("profilometry_points", (coupon_id, row["grid_x"], row["grid_y"], row["thickness_um"]))
		for row in self.reader.dicts(f"{rel_dir}/nanoindentation.csv"):
			self.put("nanoindentation", (
				coupon_id, row["indent_id"], row["x_um"], row["y_um"],
				row["hardness_gpa"], row["reduced_modulus_gpa"], row["max_load_mn"],
			))
		afm = self.reader.dicts(f"{rel_dir}/afm.csv")[0]
		self.put("afm", (coupon_id, afm["ra_nm"], afm["rq_nm"], afm["image_file"]))

		if is_friction:
			for t in range(1, fa["tracks_per_friction_coupon"] + 1):
				track_id += 1
				self.load_track(f"{rel_dir}/track_T{t}", track_id, coupon_id)
		return track_id

	def load_track(self, rel_dir: str, track_id: int, coupon_id: int) -> None:
		info = self.reader.dicts(f"{rel_dir}/track_info.csv")[0]
		self.put("tracks", (
			track_id, info["track_code"], coupon_id, self.run_ids[info["run_code"]],
			info["environment"], info["load_mn"], info["stroke_mm"], info["speed_mm_s"],
			info["counterface_id"], info["started_at"],
		))
		self.flush("tracks")

		cofs: list[float] = []
		for row in self.reader.dicts(f"{rel_dir}/cof_vs_cycle.csv"):
			cofs.append(row["cof"])
			self.put("friction_cycles", (track_id, row["cycle"], row["cof"]))
		ss_mean, ss_std, run_in = summarize_track(cofs)
		self.put("track_summary", (track_id, ss_mean, ss_std, run_in))

		pt = 0
		last_cycle = None
		for row in self.reader.dicts(f"{rel_dir}/friction_loops.csv"):
			pt = pt + 1 if row["cycle"] == last_cycle else 1
			last_cycle = row["cycle"]
			self.put("friction_loop_points", (track_id, row["cycle"], pt, row["position_um"], row["friction_force_mn"]))

		wear = self.reader.dicts(f"{rel_dir}/wear.csv")[0]
		self.put("wear", (track_id, wear["wear_volume_um3"], wear["k_archard"], wear["sliding_distance_m"]))

	# --- post-load ---

	def finalize(self) -> None:
		for ddl in INDEX_DDL:
			self.conn.execute(ddl)
		self.conn.commit()
		self.conn.execute("PRAGMA journal_mode=WAL")
		self.conn.execute("ANALYZE")
		self.conn.execute("VACUUM")
		logger.info("indexes created, WAL enabled, ANALYZE + VACUUM done")

	def validate(self) -> dict[str, int]:
		cfg = self.cfg
		v = cfg["volumes"]
		expected = {
			"instruments": len(cfg["instruments"]),
			"batches": cfg["hierarchy"]["batches"],
			"runs": cfg["friction_assignment"]["runs"],
			"wafers": cfg["hierarchy"]["batches"] * cfg["hierarchy"]["wafers_per_batch"],
			"coupons": v["coupons_total"],
			"tracks": v["tracks_total"],
			"simtra_profiles": cfg["hierarchy"]["batches"] * v["simtra_rows_per_batch"],
			"xrf_points": v["xrf_points_total"],
			"profilometry_points": v["coupons_total"] * v["profilometry_points_per_coupon"],
			"nanoindentation": v["coupons_total"] * v["nanoindentation_indents_per_coupon"],
			"afm": v["coupons_total"],
			"friction_cycles": v["cycle_rows_total"],
			"friction_loop_points": v["loop_points_total"],
			"wear": v["tracks_total"],
			"track_summary": v["tracks_total"],
		}
		counts: dict[str, int] = {}
		for table, exp in expected.items():
			got = self.conn.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0]
			counts[table] = got
			if got != exp:
				raise ValidationError(f"row count mismatch: {table}: loaded {got} != expected {exp}")
		reserve = self.conn.execute("SELECT COUNT(*) FROM coupons WHERE run_id IS NULL").fetchone()[0]
		if reserve != self.cfg["friction_assignment"]["reserve_coupons_total"]:
			raise ValidationError(f"reserve coupons: {reserve} != expected")
		fk_violations = self.conn.execute("PRAGMA foreign_key_check").fetchall()
		if fk_violations:
			raise ValidationError(f"foreign_key_check reported {len(fk_violations)} violations")
		logger.info("all %d table counts match, %d reserve coupons, foreign_key_check clean", len(expected), reserve)
		return counts


def main() -> int:
	args = parse_task_args("Task 05: convert the CSV corpus to SQLite")
	out_root: Path = args.out_root
	db_path = out_root / "sqlite" / "tribo.db"
	try:
		check_dependencies(out_root, DEPENDS_ON)
		remove_stale_marker(out_root, TASK_ID)
		for stale in (db_path, db_path.with_suffix(".db-wal"), db_path.with_suffix(".db-shm")):
			if stale.exists():
				logger.info("re-run: removing %s", stale)
				stale.unlink()
		cfg = load_lab_config(out_root)

		loader = SqliteLoader(cfg, out_root / "csv", db_path)
		loader.load_flat()
		loader.load_hierarchy()
		loader.finalize()
		counts = loader.validate()
		loader.conn.close()

		db_bytes = db_path.stat().st_size
		update_storage_sizes(out_root / "bench", "sqlite", [("db", db_bytes)])
		marker_path = write_marker(out_root, TASK_ID, {
			"db_file": db_path.as_posix(),
			"db_bytes": db_bytes,
			"tables": len(counts),
			"total_rows": sum(counts.values()),
			"friction_cycles": counts["friction_cycles"],
			"friction_loop_points": counts["friction_loop_points"],
			"tracks": counts["tracks"],
			"track_summary": counts["track_summary"],
		})
	except Exception:
		logger.critical("task %s failed", TASK_ID, exc_info=True)
		return 1

	print(f"task 05 ok: {db_path} ({db_bytes / (1024 * 1024):.1f} MiB, {len(counts)} tables, {sum(counts.values())} rows)")
	print(f"bulk rows: cycles={counts['friction_cycles']} loop_points={counts['friction_loop_points']} "
		f"xrf={counts['xrf_points']} track_summary={counts['track_summary']}")
	print(f"marker: {marker_path}")
	return 0


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