"""Task 06 - Convert the CSV corpus to PostgreSQL.

Loads the same logical schema as task 05 into the PostgreSQL 16 database
named by the TRIBO_PG_DSN environment variable: COPY FROM STDIN, RANGE
partitioning of the two bulk tables by track_id (12 partitions of 120
tracks), constraints and indexes created after the load (variant B:
enforced foreign keys everywhere - docs/research/Tribology_FK_Architecture.html),
track_summary as a materialized view (algorithm fixed in
docs/rules/db-sql-schema.md section 5, cross-checked against
common/track_summary.py). Host CPU/RAM during the load is sampled from the
database host's Prometheus (TRIBO_PROM_URL). Appends live (and, when
pg_dump is available, dump) sizes to storage_sizes.csv and writes
./out/.done/06.ok. Spec: docs/specs/06_convert_postgresql.md.
"""

from __future__ import annotations

import logging
import shutil
import subprocess
import sys
import time
from pathlib import Path

import psycopg

from common.corpus import CorpusReader, ValidationError
from common.monitoring import get_prom_url, host_window_metrics
from common.pg import get_dsn
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
from common.track_summary import summarize_track

logger = logging.getLogger(__name__)

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

BATCH_ROWS = 50000  # rows per COPY chunk (mirrors the 50k batching of spec 05)
PARTITIONS = 12  # RANGE partitions of 120 tracks each, spec 06
TRACKS_PER_PARTITION = 120
SUMMARY_PARITY_TRACKS = (1, 720, 1440)  # matview cross-checked vs common/track_summary.py
PARITY_TOLERANCE = 1e-9  # cross-engine float tolerance, test-pipeline-validation

# Float columns are DOUBLE PRECISION, not REAL: the corpus is float64 and
# REAL (float4) would break the 1e-9 cross-engine result parity that the
# whole benchmark depends on (db-sql-schema section 6; spec updated).
CREATE_TABLES = [
	"CREATE TABLE instruments (instrument_id INTEGER NOT NULL, instrument_code TEXT NOT NULL, name TEXT NOT NULL, role TEXT NOT NULL)",
	"CREATE TABLE batches (batch_id INTEGER NOT NULL, batch_code TEXT NOT NULL, 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 DATE NOT NULL)",
	"CREATE TABLE runs (run_id INTEGER NOT NULL, run_code TEXT NOT NULL, environment TEXT NOT NULL, date DATE NOT NULL, plates INTEGER NOT NULL, operator TEXT NOT NULL)",
	"CREATE TABLE wafers (wafer_id INTEGER NOT NULL, wafer_code TEXT NOT NULL, batch_id INTEGER NOT NULL, wafer_index INTEGER NOT NULL, deposition_date DATE NOT NULL, coupons INTEGER NOT NULL, friction_coupons INTEGER NOT NULL, reserve_coupons INTEGER NOT NULL)",
	"CREATE TABLE coupons (coupon_id INTEGER NOT NULL, coupon_code TEXT NOT NULL, wafer_id INTEGER NOT NULL, batch_id INTEGER NOT NULL, grid_row INTEGER NOT NULL, grid_col INTEGER NOT NULL, thickness_um DOUBLE PRECISION NOT NULL, ra_nm DOUBLE PRECISION NOT NULL, au_wtpct_mean DOUBLE PRECISION NOT NULL, run_id INTEGER, plate INTEGER, probe INTEGER, square INTEGER)",
	"CREATE TABLE tracks (track_id INTEGER NOT NULL, track_code TEXT NOT NULL, coupon_id INTEGER NOT NULL, run_id INTEGER NOT NULL, environment TEXT NOT NULL, load_mn DOUBLE PRECISION NOT NULL, stroke_mm DOUBLE PRECISION NOT NULL, speed_mm_s DOUBLE PRECISION NOT NULL, counterface_id TEXT NOT NULL, started_at TIMESTAMPTZ NOT NULL)",
	"CREATE TABLE simtra_profiles (batch_id INTEGER NOT NULL, row_no INTEGER NOT NULL, angle_deg DOUBLE PRECISION NOT NULL, energy_ev DOUBLE PRECISION NOT NULL, pt_flux DOUBLE PRECISION NOT NULL, au_flux DOUBLE PRECISION NOT NULL)",
	"CREATE TABLE xrf_points (coupon_id INTEGER NOT NULL, grid_x INTEGER NOT NULL, grid_y INTEGER NOT NULL, pt_wtpct DOUBLE PRECISION NOT NULL, au_wtpct DOUBLE PRECISION NOT NULL)",
	"CREATE TABLE profilometry_points (coupon_id INTEGER NOT NULL, grid_x INTEGER NOT NULL, grid_y INTEGER NOT NULL, thickness_um DOUBLE PRECISION NOT NULL)",
	"CREATE TABLE nanoindentation (coupon_id INTEGER NOT NULL, indent_id INTEGER NOT NULL, x_um DOUBLE PRECISION NOT NULL, y_um DOUBLE PRECISION NOT NULL, hardness_gpa DOUBLE PRECISION NOT NULL, reduced_modulus_gpa DOUBLE PRECISION NOT NULL, max_load_mn DOUBLE PRECISION NOT NULL)",
	"CREATE TABLE afm (coupon_id INTEGER NOT NULL, ra_nm DOUBLE PRECISION NOT NULL, rq_nm DOUBLE PRECISION NOT NULL, image_file TEXT NOT NULL)",
	"CREATE TABLE friction_cycles (track_id INTEGER NOT NULL, cycle INTEGER NOT NULL, cof DOUBLE PRECISION NOT NULL) PARTITION BY RANGE (track_id)",
	"CREATE TABLE friction_loop_points (track_id INTEGER NOT NULL, cycle INTEGER NOT NULL, pt INTEGER NOT NULL, position_um DOUBLE PRECISION NOT NULL, friction_force_mn DOUBLE PRECISION NOT NULL) PARTITION BY RANGE (track_id)",
	"CREATE TABLE wear (track_id INTEGER NOT NULL, wear_volume_um3 DOUBLE PRECISION NOT NULL, k_archard DOUBLE PRECISION NOT NULL, sliding_distance_m DOUBLE PRECISION NOT NULL)",
]

CONSTRAINTS = [
	"ALTER TABLE instruments ADD PRIMARY KEY (instrument_id)",
	"ALTER TABLE instruments ADD UNIQUE (instrument_code)",
	"ALTER TABLE batches ADD PRIMARY KEY (batch_id)",
	"ALTER TABLE batches ADD UNIQUE (batch_code)",
	"ALTER TABLE runs ADD PRIMARY KEY (run_id)",
	"ALTER TABLE runs ADD UNIQUE (run_code)",
	"ALTER TABLE wafers ADD PRIMARY KEY (wafer_id)",
	"ALTER TABLE wafers ADD UNIQUE (wafer_code)",
	"ALTER TABLE coupons ADD PRIMARY KEY (coupon_id)",
	"ALTER TABLE coupons ADD UNIQUE (coupon_code)",
	"ALTER TABLE tracks ADD PRIMARY KEY (track_id)",
	"ALTER TABLE tracks ADD UNIQUE (track_code)",
	"ALTER TABLE simtra_profiles ADD PRIMARY KEY (batch_id, row_no)",
	"ALTER TABLE xrf_points ADD PRIMARY KEY (coupon_id, grid_x, grid_y)",
	"ALTER TABLE profilometry_points ADD PRIMARY KEY (coupon_id, grid_x, grid_y)",
	"ALTER TABLE nanoindentation ADD PRIMARY KEY (coupon_id, indent_id)",
	"ALTER TABLE afm ADD PRIMARY KEY (coupon_id)",
	"ALTER TABLE friction_cycles ADD PRIMARY KEY (track_id, cycle)",
	"ALTER TABLE friction_loop_points ADD PRIMARY KEY (track_id, cycle, pt)",
	"ALTER TABLE wear ADD PRIMARY KEY (track_id)",
	"ALTER TABLE wafers ADD FOREIGN KEY (batch_id) REFERENCES batches",
	"ALTER TABLE coupons ADD FOREIGN KEY (wafer_id) REFERENCES wafers",
	"ALTER TABLE coupons ADD FOREIGN KEY (batch_id) REFERENCES batches",
	"ALTER TABLE coupons ADD FOREIGN KEY (run_id) REFERENCES runs",
	"ALTER TABLE tracks ADD FOREIGN KEY (coupon_id) REFERENCES coupons",
	"ALTER TABLE tracks ADD FOREIGN KEY (run_id) REFERENCES runs",
	"ALTER TABLE simtra_profiles ADD FOREIGN KEY (batch_id) REFERENCES batches",
	"ALTER TABLE xrf_points ADD FOREIGN KEY (coupon_id) REFERENCES coupons",
	"ALTER TABLE profilometry_points ADD FOREIGN KEY (coupon_id) REFERENCES coupons",
	"ALTER TABLE nanoindentation ADD FOREIGN KEY (coupon_id) REFERENCES coupons",
	"ALTER TABLE afm ADD FOREIGN KEY (coupon_id) REFERENCES coupons",
	"ALTER TABLE friction_cycles ADD FOREIGN KEY (track_id) REFERENCES tracks",
	"ALTER TABLE friction_loop_points ADD FOREIGN KEY (track_id) REFERENCES tracks",
	"ALTER TABLE wear ADD FOREIGN KEY (track_id) REFERENCES tracks",
]

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
	# Q5 pruning within partitions; cascades to every partition (spec 06).
	# nanoindentation(coupon_id) from the spec is omitted: the PK prefix serves it.
	"CREATE INDEX ix_brin_friction_cycles_cycle ON friction_cycles USING BRIN (cycle)",
]

# Materialized view semantics fixed in docs/rules/db-sql-schema.md section 5.
MATVIEW_SQL = """
CREATE MATERIALIZED VIEW track_summary AS
WITH tail AS (
	SELECT track_id, avg(cof) AS m, stddev_pop(cof) AS s
	FROM friction_cycles
	WHERE cycle >= 501
	GROUP BY track_id
),
run_in AS (
	SELECT fc.track_id,
	       COALESCE(MIN(fc.cycle) FILTER (WHERE abs(fc.cof - t.m) < 2 * t.s), 500) AS run_in_cycles
	FROM friction_cycles fc
	JOIN tail t USING (track_id)
	GROUP BY fc.track_id
)
SELECT r.track_id,
       avg(fc.cof) AS cof_ss_mean,
       stddev_pop(fc.cof) AS cof_ss_std,
       r.run_in_cycles
FROM run_in r
JOIN friction_cycles fc ON fc.track_id = r.track_id AND fc.cycle > r.run_in_cycles
GROUP BY r.track_id, r.run_in_cycles
"""


def partition_ddl() -> list[str]:
	ddl = []
	for parent in ("friction_cycles", "friction_loop_points"):
		for i in range(PARTITIONS):
			lo = i * TRACKS_PER_PARTITION + 1
			hi = lo + TRACKS_PER_PARTITION
			ddl.append(f"CREATE TABLE {parent}_p{i + 1:02d} PARTITION OF {parent} FOR VALUES FROM ({lo}) TO ({hi})")
	return ddl


class PgLoader:
	def __init__(self, conn: psycopg.Connection) -> None:
		self.conn = conn
		self.buffers: dict[str, list[tuple]] = {t: [] for t in TABLES}
		self.summary_parity: dict[int, tuple] = {}

	def create_schema(self) -> None:
		with self.conn.cursor() as cur:
			cur.execute("DROP MATERIALIZED VIEW IF EXISTS track_summary")
			for table in reversed(TABLES):
				if table != "track_summary":
					cur.execute(f"DROP TABLE IF EXISTS {table} CASCADE")
			for ddl in CREATE_TABLES + partition_ddl():
				cur.execute(ddl)
		self.conn.commit()
		logger.info("schema created: %d tables, %d partitions", len(CREATE_TABLES), 2 * PARTITIONS)

	def load(self, cfg: dict, reader: CorpusReader) -> None:
		pending = 0
		for table, row in stream_rows(cfg, reader):
			if table == "track_summary":
				# computed server-side as a materialized view; keep a sample
				# of the Python values for the cross-engine parity check
				if row[0] in SUMMARY_PARITY_TRACKS:
					self.summary_parity[row[0]] = row
				continue
			self.buffers[table].append(row)
			pending += 1
			if pending >= BATCH_ROWS:
				self.flush_all()
				pending = 0
		self.flush_all()

	def flush_all(self) -> None:
		for table in TABLES:
			buf = self.buffers.get(table)
			if not buf:
				continue
			cols = ", ".join(COLUMNS[table])
			with self.conn.cursor() as cur:
				with cur.copy(f"COPY {table} ({cols}) FROM STDIN") as copy:
					for row in buf:
						copy.write_row(row)
			buf.clear()
		self.conn.commit()

	def finalize(self) -> None:
		with self.conn.cursor() as cur:
			for ddl in CONSTRAINTS:
				cur.execute(ddl)
			for ddl in INDEX_DDL:
				cur.execute(ddl)
			cur.execute(MATVIEW_SQL)
			cur.execute("CREATE UNIQUE INDEX ix_track_summary_track_id ON track_summary(track_id)")  # Q2/Q3 joins
		self.conn.commit()
		logger.info("constraints, indexes and track_summary materialized view created")

	def validate(self, cfg: dict) -> dict[str, int]:
		counts: dict[str, int] = {}
		with self.conn.cursor() as cur:
			for table, exp in expected_counts(cfg).items():
				got = cur.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 = cur.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")
			for track_id, py_row in self.summary_parity.items():
				db_row = cur.execute(
					"SELECT cof_ss_mean, cof_ss_std, run_in_cycles FROM track_summary WHERE track_id = %s",
					(track_id,),
				).fetchone()
				if db_row is None:
					raise ValidationError(f"track_summary missing track {track_id}")
				if (abs(db_row[0] - py_row[1]) > PARITY_TOLERANCE
						or abs(db_row[1] - py_row[2]) > PARITY_TOLERANCE
						or db_row[2] != py_row[3]):
					raise ValidationError(
						f"track_summary parity failed for track {track_id}: db={db_row} py={py_row[1:]}"
					)
		logger.info(
			"all %d table counts match, %d reserve coupons, matview parity ok for tracks %s",
			len(counts), reserve, list(self.summary_parity),
		)
		return counts

	def table_sizes(self) -> dict[str, int]:
		sizes: dict[str, int] = {}
		with self.conn.cursor() as cur:
			for table in TABLES:
				total = cur.execute(
					"""SELECT COALESCE(pg_total_relation_size(%s::regclass), 0)
					 + COALESCE((SELECT sum(pg_total_relation_size(inhrelid))
					             FROM pg_inherits WHERE inhparent = %s::regclass), 0)""",
					(table, table),
				).fetchone()[0]
				sizes[table] = int(total)
		return sizes


def make_dump(dsn: str, dump_path: Path) -> int:
	pg_dump = shutil.which("pg_dump")
	if not pg_dump:
		logger.warning("pg_dump skipped (client not installed on this host; produce the dump on the database host)")
		return 0
	dump_path.parent.mkdir(parents=True, exist_ok=True)
	subprocess.run([pg_dump, "--format=custom", f"--file={dump_path}", dsn], check=True)
	return dump_path.stat().st_size


def main() -> int:
	args = parse_task_args("Task 06: convert the CSV corpus to PostgreSQL")
	out_root: Path = args.out_root
	dsn = get_dsn()
	try:
		check_dependencies(out_root, DEPENDS_ON)
		remove_stale_marker(out_root, TASK_ID)
		cfg = load_lab_config(out_root)

		t_start = time.time()
		with psycopg.connect(dsn) as conn:
			database = conn.info.dbname
			logger.info("connected to database %r on %s", database, conn.info.host)
			loader = PgLoader(conn)
			loader.create_schema()
			loader.load(cfg, CorpusReader(out_root / "csv"))
			load_done = time.time()
			loader.finalize()
			counts = loader.validate(cfg)
			sizes = loader.table_sizes()
		with psycopg.connect(dsn, autocommit=True) as conn:
			conn.execute("VACUUM ANALYZE")
		t_end = time.time()

		live_bytes = sum(sizes.values())
		variants = [("live", live_bytes)]
		dump_bytes = make_dump(dsn, out_root / "pg" / "tribo.dump")
		if dump_bytes:
			variants.append(("dump", dump_bytes))
		update_storage_sizes(out_root / "bench", "pg", variants)

		entries = {
			"database": database,
			"tables": len(counts),
			"total_rows": sum(counts.values()),
			"live_bytes": live_bytes,
			"friction_cycles_bytes": sizes["friction_cycles"],
			"friction_loop_points_bytes": sizes["friction_loop_points"],
			"copy_seconds": round(load_done - t_start, 1),
			"total_seconds": round(t_end - t_start, 1),
			"dump": dump_bytes if dump_bytes else "skipped (pg_dump unavailable on this host)",
		}
		entries.update(process_metrics())  # client-side converter cost
		prom_url = get_prom_url()
		if prom_url:
			metrics = host_window_metrics(prom_url, t_start, t_end)
			if metrics:
				entries.update(metrics)
				logger.info("db host during load: %s", 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 06 ok: database {database} ({live_bytes / (1024 * 1024):.1f} MiB live, "
		f"{len(counts)} tables, {sum(counts.values())} rows, {entries['total_seconds']} s)")
	print(f"bulk rows: cycles={counts['friction_cycles']} loop_points={counts['friction_loop_points']} "
		f"track_summary={counts['track_summary']}")
	if prom_url and "host_cpu_avg_pct" in entries:
		print(f"db host: cpu avg {entries['host_cpu_avg_pct']}% max {entries['host_cpu_max_pct']}%, "
			f"ram peak {entries['host_ram_used_peak_mb']} MB (delta {entries['host_ram_used_delta_mb']} MB)")
	print(f"marker: {marker_path}")
	return 0


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