from __future__ import annotations

import json
import math
import sys
import argparse
from pathlib import Path

try:
    import yaml
except ImportError:
    raise SystemExit("Install dependencies: python -m pip install -r requirements.txt")

try:
    from pint import DimensionalityError, UndefinedUnitError, UnitRegistry
except ImportError:
    raise SystemExit("Install dependencies: python -m pip install -r requirements.txt")


HERE = Path(__file__).resolve().parent

ureg = UnitRegistry()
ureg.define("kip = 1000 * force_pound")
ureg.define("ksi = kip / inch ** 2")
ureg.define("psf = force_pound / foot ** 2")
ureg.define("plf = force_pound / foot")  # pounds per linear foot
if "psi" not in ureg:
    ureg.define("psi = force_pound / inch ** 2")


def quantity(value, unit: str, name: str) -> float:
    try:
        q = ureg.Quantity(value).to(unit)
    except (DimensionalityError, UndefinedUnitError, TypeError, ValueError) as exc:
        raise ValueError(f"{name}: expected {unit}, got {value!r}") from exc
    magnitude = float(q.magnitude)
    if magnitude <= 0:
        raise ValueError(f"{name} must be positive")
    return magnitude


def factor(value, name: str) -> float:
    magnitude = float(value)
    if magnitude <= 0:
        raise ValueError(f"{name} must be positive")
    return magnitude


def compute(inp: dict) -> dict:
    span_ft  = quantity(inp["span"], "ft", "span")
    b_in     = quantity(inp["width"], "in", "width")
    d_in     = quantity(inp["depth"], "in", "depth")
    lb_in    = quantity(inp["bearing_length"], "in", "bearing_length")
    bed_in   = quantity(inp["bearing_end_distance"], "in", "bearing_end_distance")
    le_in    = quantity(inp["unbraced_length"], "in", "unbraced_length")
    Fb_psi   = quantity(inp["Fb"], "psi", "Fb")
    Fv_psi   = quantity(inp["Fv"], "psi", "Fv")
    Fcp_psi  = quantity(inp["Fcp"], "psi", "Fcp")
    E_psi    = quantity(inp["E"], "psi", "E")
    Emin_psi = quantity(inp["E_min"], "psi", "E_min")
    species_grade = str(inp["species_grade"])
    CD = factor(inp["CD"], "CD")
    CM = factor(inp["CM"], "CM")
    Ct = factor(inp["Ct"], "Ct")
    CF = factor(inp["CF"], "CF")
    Cr = factor(inp["Cr"], "Cr")
    Cfu = factor(inp["Cfu"], "Cfu")
    Ci = factor(inp["Ci"], "Ci")
    Cvr = factor(inp["Cvr"], "Cvr")

    # Checked demands supplied by YAML. beam.typ derives the same values and the
    # test suite reconciles them via the <wood-joist-loads> metadata label.
    w_plf   = quantity(inp["w"], "plf", "w")
    wL_plf  = quantity(inp["wL"], "plf", "wL")
    M_ftlbf = quantity(inp["M"], "ft * lbf", "M")
    V_lbf   = quantity(inp["V"], "lbf", "V")
    R_lbf   = quantity(inp["R"], "lbf", "R")

    A_in2  = b_in * d_in
    Ix_in4 = b_in * d_in ** 3 / 12.0
    Sx_in3 = b_in * d_in ** 2 / 6.0

    # Flexure (NDS 3.3.3)
    RB = math.sqrt(le_in * d_in / b_in ** 2)
    if RB >= 50.0:
        raise ValueError(f"RB = {RB:.3f} >= 50; member requires additional "
                         "lateral bracing (NDS 3.3.3)")
    Fb_star_psi = Fb_psi * CD * CM * Ct * CF * Ci * Cr
    FbE_psi = 1.20 * Emin_psi / RB ** 2
    ratio = FbE_psi / Fb_star_psi
    CL = (1.0 + ratio) / 1.9 - math.sqrt(((1.0 + ratio) / 1.9) ** 2 - ratio / 0.95)
    if not math.isfinite(CL) or CL > 1.0 or CL <= 0:
        raise ValueError(f"CL = {CL:.4f} out of range (0, 1]; beam stability "
                         "factor invalid (NDS 3.3.3)")
    Fb_prime_psi = Fb_star_psi * CL * Cfu
    fb_psi = M_ftlbf * 12.0 / Sx_in3
    Ma_ftlbf = Sx_in3 * Fb_prime_psi / 12.0

    # Shear (NDS 3.4)
    Fv_prime_psi = Fv_psi * CD * CM * Ct * Cvr
    fv_psi = 1.5 * V_lbf / A_in2

    # Bearing (NDS 3.10)
    if lb_in >= 6.0:
        raise ValueError("bearing_length must be < 6 in for the bearing area "
                         "factor (NDS 3.10.4)")
    if bed_in < 3.0:
        raise ValueError("bearing_end_distance must be >= 3 in for the bearing "
                         "area factor (NDS 3.10.4)")
    Cb = (lb_in + 0.375) / lb_in
    Fcp_prime_psi = Fcp_psi * CM * Ct * Ci * Cb
    Ab_in2 = lb_in * b_in
    fcp_psi = R_lbf / Ab_in2

    # Serviceability (NDS 3.5). Limits hardcoded per reference: L/240, L/360.
    L_in = span_ft * 12.0
    delta_total_in = 5.0 * (w_plf / 12.0) * L_in ** 4 / (384.0 * E_psi * Ix_in4)
    delta_live_in = 5.0 * (wL_plf / 12.0) * L_in ** 4 / (384.0 * E_psi * Ix_in4)
    delta_allow_total_in = L_in / 240.0
    delta_allow_live_in = L_in / 360.0

    def q(value: float) -> float:
        return round(value, 6)

    values = {
        "span_ft": q(span_ft),
        "w_plf": q(w_plf), "wL_plf": q(wL_plf),
        "b_in": q(b_in), "d_in": q(d_in),
        "A_in2": q(A_in2), "Ix_in4": q(Ix_in4), "Sx_in3": q(Sx_in3),
        "M_ftlbf": q(M_ftlbf), "V_lbf": q(V_lbf), "R_lbf": q(R_lbf),
        "species_grade": species_grade,
        "Fb_psi": q(Fb_psi), "Fv_psi": q(Fv_psi), "Fcp_psi": q(Fcp_psi),
        "E_psi": q(E_psi), "Emin_psi": q(Emin_psi),
        "CD": q(CD), "CM": q(CM), "Ct": q(Ct), "CF": q(CF),
        "Cr": q(Cr), "Cfu": q(Cfu), "Ci": q(Ci), "Cvr": q(Cvr),
        "le_in": q(le_in), "RB": q(RB),
        "Fb_star_psi": q(Fb_star_psi), "FbE_psi": q(FbE_psi),
        "CL": q(CL), "Fb_prime_psi": q(Fb_prime_psi),
        "fb_psi": q(fb_psi), "Ma_ftlbf": q(Ma_ftlbf),
        "Fv_prime_psi": q(Fv_prime_psi), "fv_psi": q(fv_psi),
        "lb_in": q(lb_in), "bed_in": q(bed_in), "Cb": q(Cb), "Ab_in2": q(Ab_in2),
        "Fcp_prime_psi": q(Fcp_prime_psi), "fcp_psi": q(fcp_psi),
        "delta_total_in": q(delta_total_in), "delta_live_in": q(delta_live_in),
        "delta_allow_total_in": q(delta_allow_total_in),
        "delta_allow_live_in": q(delta_allow_live_in),
    }

    checks = {
        "flexure":          {"demand": q(fb_psi), "capacity": q(Fb_prime_psi), "ok": fb_psi <= Fb_prime_psi},
        "shear":            {"demand": q(fv_psi), "capacity": q(Fv_prime_psi), "ok": fv_psi <= Fv_prime_psi},
        "bearing":          {"demand": q(fcp_psi), "capacity": q(Fcp_prime_psi), "ok": fcp_psi <= Fcp_prime_psi},
        "deflection_total": {"demand": q(delta_total_in), "capacity": q(delta_allow_total_in), "ok": delta_total_in <= delta_allow_total_in},
        "deflection_live":  {"demand": q(delta_live_in), "capacity": q(delta_allow_live_in), "ok": delta_live_in <= delta_allow_live_in},
    }

    return {
        "tool": "wood_joist",
        "version": "0.1",
        "project": inp.get("project", ""),
        "prepared_by": inp.get("prepared_by", ""),
        "values": values,
        "checks": checks,
    }


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Calculate the wood joist design from a YAML input file.")
    parser.add_argument("--input", "-i", type=Path, default=HERE / "input.yaml", help="YAML input path")
    parser.add_argument("--output", "-o", type=Path, help="JSON output path; defaults beside the input")
    parser.add_argument("--stdout", action="store_true", help="Write the complete JSON result to stdout instead of a file")
    args = parser.parse_args(argv)

    input_path = args.input
    output_path = args.output or input_path.with_name("results.json")
    with input_path.open(encoding="utf-8") as handle:
        result = compute(yaml.safe_load(handle))
    serialized = json.dumps(result, indent=2) + "\n"
    if args.stdout:
        sys.stdout.write(serialized)
    else:
        output_path.write_text(serialized, encoding="utf-8")
        print(output_path)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())