"""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,
	process_metrics,
	remove_stale_marker,
	write_marker,
)
from common.relational import COLUMNS, TABLES, expected_counts, stream_rows
from common.storage_sizes import update_storage_sizes

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
]


class SqliteLoader:
	def __init__(self, db_path: Path) -> None:
		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.inserts = {
			t: f"INSERT INTO {t} VALUES ({','.join('?' * len(COLUMNS[t]))})" for t in TABLES
		}
		self.buffers: dict[str, list[tuple]] = {t: [] for t in TABLES}

	def load(self, cfg: dict, reader: CorpusReader) -> None:
		pending = 0
		for table, row in stream_rows(cfg, reader):
			self.buffers[table].append(row)
			pending += 1
			if pending >= BATCH_ROWS:
				self.flush_all()
				pending = 0
		self.flush_all()

	def flush_all(self) -> None:
		# TABLES is FK-dependency ordered: parents flush before children.
		for table in TABLES:
			buf = self.buffers[table]
			if buf:
				self.conn.executemany(self.inserts[table], buf)
				buf.clear()
		self.conn.commit()

	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, cfg: dict) -> dict[str, int]:
		counts: dict[str, int] = {}
		for table, exp in expected_counts(cfg).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 != 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(counts), 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(db_path)
		loader.load(cfg, CorpusReader(out_root / "csv"))
		loader.finalize()
		counts = loader.validate(cfg)
		loader.conn.close()

		db_bytes = db_path.stat().st_size
		update_storage_sizes(out_root / "bench", "sqlite", [("db", db_bytes)])
		entries = {
			"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"],
		}
		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 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())