"""Shared relational row stream for the SQLite and PostgreSQL converters.

Walks the CSV corpus once (manifest-validated via CorpusReader) and yields
(table, row_tuple) pairs in FK-dependency-safe order: a parent row is always
yielded before any row that references it. Both engines load the SAME
logical schema from this stream (docs/rules/db-sql-schema.md section 6).
"""

from __future__ import annotations

import logging
from typing import Iterator

from common.corpus import CorpusReader
from common.track_summary import summarize_track

logger = logging.getLogger(__name__)

# Dependency-safe order; consumers that buffer rows must flush in this order.
TABLES = [
	"instruments",
	"batches",
	"runs",
	"wafers",
	"coupons",
	"tracks",
	"simtra_profiles",
	"xrf_points",
	"profilometry_points",
	"nanoindentation",
	"afm",
	"friction_cycles",
	"friction_loop_points",
	"wear",
	"track_summary",
]

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

EXPECTED_COUNTS_KEYS = TABLES  # every table is count-validated after load


def expected_counts(cfg: dict) -> dict[str, int]:
	v = cfg["volumes"]
	h = cfg["hierarchy"]
	return {
		"instruments": len(cfg["instruments"]),
		"batches": h["batches"],
		"runs": cfg["friction_assignment"]["runs"],
		"wafers": h["batches"] * h["wafers_per_batch"],
		"coupons": v["coupons_total"],
		"tracks": v["tracks_total"],
		"simtra_profiles": h["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"],
	}


def stream_rows(cfg: dict, reader: CorpusReader) -> Iterator[tuple[str, tuple]]:
	batch_ids: dict[str, int] = {}
	run_ids: dict[str, int] = {}

	for i, inst in enumerate(cfg["instruments"], 1):
		yield "instruments", (i, inst["instrument_id"], inst["name"], inst["role"])
	for i, row in enumerate(reader.dicts("batches.csv"), 1):
		batch_ids[row["batch_code"]] = i
		yield "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(reader.dicts("runs.csv"), 1):
		run_ids[row["run_code"]] = i
		yield "runs", (i, row["run_code"], row["environment"], row["date"], row["plates"], row["operator"])
	for code, batch_id in batch_ids.items():
		for row_no, row in enumerate(reader.dicts(f"simtra/simtra_profile_{code}.csv"), 1):
			yield "simtra_profiles", (batch_id, row_no, row["angle_deg"], row["energy_ev"], row["pt_flux"], row["au_flux"])

	h = cfg["hierarchy"]
	fa = cfg["friction_assignment"]
	wafer_id = coupon_id = track_id = 0
	for batch in cfg["deposition_matrix"]:
		b_code = batch["batch_code"]
		for w in range(1, h["wafers_per_batch"] + 1):
			wafer_id += 1
			wafer_dir = f"batch_{b_code}/wafer_W{w}"
			info = reader.dicts(f"{wafer_dir}/wafer_info.csv")[0]
			yield "wafers", (
				wafer_id, info["wafer_code"], batch_ids[b_code], info["wafer_index"],
				info["deposition_date"], info["coupons"], info["friction_coupons"], info["reserve_coupons"],
			)
			for c in range(1, h["coupons_per_wafer"] + 1):
				coupon_id += 1
				rel_dir = f"{wafer_dir}/coupon_C{c:02d}"
				info = reader.dicts(f"{rel_dir}/coupon_info.csv")[0]
				is_friction = info["run_code"] != "RESERVE"
				yield "coupons", (
					coupon_id, info["coupon_code"], wafer_id, batch_ids[b_code],
					info["grid_row"], info["grid_col"], info["thickness_um"], info["ra_nm"],
					info["au_wtpct_mean"], run_ids[info["run_code"]] if is_friction else None,
					info.get("plate"), info.get("probe"), info.get("square"),
				)
				for row in reader.dicts(f"{rel_dir}/xrf_map.csv"):
					yield "xrf_points", (coupon_id, row["grid_x"], row["grid_y"], row["pt_wtpct"], row["au_wtpct"])
				for row in reader.dicts(f"{rel_dir}/profilometry.csv"):
					yield "profilometry_points", (coupon_id, row["grid_x"], row["grid_y"], row["thickness_um"])
				for row in reader.dicts(f"{rel_dir}/nanoindentation.csv"):
					yield "nanoindentation", (
						coupon_id, row["indent_id"], row["x_um"], row["y_um"],
						row["hardness_gpa"], row["reduced_modulus_gpa"], row["max_load_mn"],
					)
				afm = reader.dicts(f"{rel_dir}/afm.csv")[0]
				yield "afm", (coupon_id, afm["ra_nm"], afm["rq_nm"], afm["image_file"])
				if not is_friction:
					continue
				for t in range(1, fa["tracks_per_friction_coupon"] + 1):
					track_id += 1
					track_dir = f"{rel_dir}/track_T{t}"
					tinfo = reader.dicts(f"{track_dir}/track_info.csv")[0]
					yield "tracks", (
						track_id, tinfo["track_code"], coupon_id, run_ids[tinfo["run_code"]],
						tinfo["environment"], tinfo["load_mn"], tinfo["stroke_mm"], tinfo["speed_mm_s"],
						tinfo["counterface_id"], tinfo["started_at"],
					)
					cofs: list[float] = []
					for row in reader.dicts(f"{track_dir}/cof_vs_cycle.csv"):
						cofs.append(row["cof"])
						yield "friction_cycles", (track_id, row["cycle"], row["cof"])
					ss_mean, ss_std, run_in = summarize_track(cofs)
					yield "track_summary", (track_id, ss_mean, ss_std, run_in)
					pt = 0
					last_cycle = None
					for row in reader.dicts(f"{track_dir}/friction_loops.csv"):
						pt = pt + 1 if row["cycle"] == last_cycle else 1
						last_cycle = row["cycle"]
						yield "friction_loop_points", (track_id, row["cycle"], pt, row["position_um"], row["friction_force_mn"])
					wear = reader.dicts(f"{track_dir}/wear.csv")[0]
					yield "wear", (track_id, wear["wear_volume_um3"], wear["k_archard"], wear["sliding_distance_m"])
		logger.info("batch %s streamed (through coupon %d, track %d)", b_code, coupon_id, track_id)
