from __future__ import annotations

import json
import math
import sys
import argparse
from pathlib import Path
from zipfile import ZipFile
from xml.etree import ElementTree as ET

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
DATABASE = HERE / "aisc-shapes-database-v15.0.xlsx"
NS = {"m": "http://schemas.openxmlformats.org/spreadsheetml/2006/main"}

ureg = UnitRegistry()
ureg.define("kip = 1000 * force_pound")
ureg.define("ksi = kip / inch ** 2")
ureg.define("psf = force_pound / foot ** 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 _shared_strings(book: ZipFile) -> list[str]:
    root = ET.fromstring(book.read("xl/sharedStrings.xml"))
    return ["".join(t.text or "" for t in item.findall(".//m:t", NS)) for item in root.findall("m:si", NS)]


def _cell_value(cell: ET.Element, shared: list[str]):
    value = cell.find("m:v", NS)
    if value is None:
        return None
    text = value.text or ""
    if cell.get("t") == "s":
        return shared[int(text)]
    try:
        return float(text)
    except ValueError:
        return text


def load_section(label: str, path: Path = DATABASE) -> dict[str, float | str]:
    with ZipFile(path) as book:
        shared = _shared_strings(book)
        root = ET.fromstring(book.read("xl/worksheets/sheet2.xml"))
    rows = root.findall(".//m:sheetData/m:row", NS)
    headers: dict[str, int] = {}
    for cell in rows[0].findall("m:c", NS):
        # The workbook repeats the metric section after the first 84 columns;
        # the first occurrence is the US customary section used here.
        headers.setdefault(str(_cell_value(cell, shared)), _column_number(cell.get("r", "A1")))
    label_column = headers["AISC_Manual_Label"]
    wanted = {
        # Use detailing dimensions for d and thicknesses, as shown in the
        # reference sheet; nominal dimensions are also retained in the source.
        "A": "A", "d": "ddet", "b": "bfdet", "tf": "tfdet", "tw": "twdet",
        "Ix": "Ix", "Zx": "Zx", "Sx": "Sx", "Iy": "Iy", "Zy": "Zy",
        "Sy": "Sy", "ry": "ry", "J": "J", "rts": "rts", "ho": "ho",
        "lambda": "h/tw",
    }
    for row in rows[1:]:
        values = {_column_number(cell.get("r", "A1")): _cell_value(cell, shared) for cell in row.findall("m:c", NS)}
        if values.get(label_column) == label:
            result: dict[str, float | str] = {"label": label}
            for name, header in wanted.items():
                raw = values.get(headers[header])
                if raw is None or raw == "–":
                    raise ValueError(f"Section {label} has no usable {header} property")
                result[name] = float(raw)
            return result
    raise ValueError(f"Section {label!r} was not found in {path.name}")


def _column_number(reference: str) -> int:
    letters = "".join(ch for ch in reference if ch.isalpha())
    number = 0
    for letter in letters:
        number = number * 26 + ord(letter.upper()) - ord("A") + 1
    return number


def compute(inp: dict, database: Path = DATABASE) -> dict:
    beam_length = quantity(inp["beam_length"], "ft", "beam_length")
    unbraced_length = quantity(inp["unbraced_length"], "in", "unbraced_length")
    moment_kipft = quantity(inp["Mu"], "kip * ft", "Mu")
    shear_kip = quantity(inp["Vu"], "kip", "Vu")
    E = quantity(inp["steel_modulus"], "ksi", "steel_modulus")
    Fy = quantity(inp["steel_yield"], "ksi", "steel_yield")
    service_load = quantity(inp["service_load"], "lbf/ft", "service_load")
    Cb = float(inp.get("cb", 1))
    c = float(inp.get("c", 1))
    if Cb <= 0 or c <= 0:
        raise ValueError("cb and c must be positive")
    section = load_section(str(inp["section"]), database)

    # The entered demands are the factored moment and the factored reaction on
    # the connector (used here as the factored shear). They arrive from the load
    # determination, which is based on a uniform gravity load, so the equivalent
    # factored uniform load is recovered from the shear: w_u = 2 V_u / L.
    factored_uniform_load_kipft = 2.0 * shear_kip / beam_length

    Lb = unbraced_length
    Lp = 1.76 * float(section["ry"]) * math.sqrt(E / Fy)
    rts = float(section["rts"])
    Sx = float(section["Sx"])
    ho = float(section["ho"])
    J = float(section["J"])
    Lr = 1.95 * rts * E / (0.7 * Fy) * math.sqrt(J * c / (Sx * ho) + math.sqrt((J * c / (Sx * ho)) ** 2 + 6.76 * (0.7 * Fy / E) ** 2))
    Fcr = Cb * math.pi**2 * E / (Lb / rts) ** 2 * math.sqrt(1 + 0.078 * J * c / (Sx * ho) * (Lb / rts) ** 2)
    Mp = Fy * float(section["Zx"]) / 12.0
    if Lb <= Lp:
        Mn_ltb = Mp
        ltb_mode = "yielding"
    elif Lb <= Lr:
        Mn_ltb = Cb * (Mp - (Mp - 0.7 * Fy * Sx / 12.0) * (Lb - Lp) / (Lr - Lp))
        ltb_mode = "inelastic LTB"
    else:
        Mn_ltb = Fcr * Sx / 12.0
        ltb_mode = "elastic LTB"
    Mn = min(Mp, Mn_ltb)
    phi_mn = 0.9 * Mn

    Aw = float(section["d"]) * float(section["tw"])
    lambda_lim = 2.24 * math.sqrt(E / Fy)
    kv = 5.34
    lambda_web = float(section["lambda"])
    cv1 = 1.0 if lambda_web <= 1.10 * math.sqrt(kv * E / Fy) else 1.10 * math.sqrt(kv * E / Fy) / lambda_web
    phi_v = 1.0 if lambda_web <= lambda_lim else 0.9
    phi_vn = phi_v * 0.6 * Fy * Aw * cv1

    # Serviceability deflection under the service uniform load (L/240 limit).
    L_in = beam_length * 12.0
    w_serv_lbf_in = service_load / 12.0
    delta_limit = L_in / 240.0
    delta = 5.0 * w_serv_lbf_in * L_in ** 4 / (384.0 * (E * 1000.0) * float(section["Ix"]))

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

    values = {
        "beam_length_ft": q(beam_length),
        "unbraced_length_in": q(Lb),
        "factored_uniform_load_kipft": q(factored_uniform_load_kipft),
        "service_load_lbf_ft": q(service_load),
        "moment_kipft": q(moment_kipft),
        "shear_kip": q(shear_kip),
        "E_ksi": q(E),
        "Fy_ksi": q(Fy),
        "Lp_in": q(Lp),
        "Lr_ft": q(Lr / 12),
        "Lb_ft": q(Lb / 12),
        "rts_in": q(rts),
        "Fcr_ksi": q(Fcr),
        "Mp_kipft": q(Mp),
        "MnLTB_kipft": q(Mn_ltb),
        "Mn_kipft": q(Mn),
        "phiMn_kipft": q(phi_mn),
        "Aw_in2": q(Aw),
        "lambda": q(lambda_web),
        "lambda_lim": q(lambda_lim),
        "kv": q(kv),
        "Cv1": q(cv1),
        "phi_v": q(phi_v),
        "phiVn_kip": q(phi_vn),
        "delta_limit_in": q(delta_limit),
        "delta_in": q(delta),
        "ltb_mode": ltb_mode,
    }
    values.update({f"section_{key}": q(float(value)) for key, value in section.items() if key != "label"})
    return {
        "tool": "steel_beam",
        "version": "0.1",
        "project": inp.get("project", ""),
        "prepared_by": inp.get("prepared_by", ""),
        "section": section["label"],
        "values": values,
        "checks": {
            "flexure": {"demand": q(moment_kipft), "capacity": q(phi_mn), "ok": moment_kipft <= phi_mn},
            "shear": {"demand": q(shear_kip), "capacity": q(phi_vn), "ok": shear_kip <= phi_vn},
            "deflection": {"demand": q(delta), "capacity": q(delta_limit), "ok": delta <= delta_limit},
        },
    }


def summary(result: dict) -> str:
    values = result["values"]
    checks = result["checks"]

    def status(name: str) -> str:
        return "OK" if checks[name]["ok"] else "NOT OK"

    return "\n".join(
        [
            "Steel Beam Design Summary",
            f"Project: {result['project']}",
            f"Section: {result['section']}",
            f"Span: {values['beam_length_ft']:.2f} ft",
            "",
            "Demands",
            f"  Factored moment, Mu: {values['moment_kipft']:.3f} kip-ft",
            f"  Factored shear, Vu:  {values['shear_kip']:.3f} kip",
            f"  Service load:        {values.get('service_load_lbf_ft', 'see input')} lbf/ft",
            "",
            "Strength",
            f"  Flexure: {status('flexure')} ({values['phiMn_kipft']:.3f} kip-ft capacity, D/C {values['moment_kipft'] / values['phiMn_kipft']:.3f})",
            f"  Shear:   {status('shear')} ({values['phiVn_kip']:.3f} kip capacity, D/C {values['shear_kip'] / values['phiVn_kip']:.3f})",
            f"  LTB mode: {values['ltb_mode']}",
            "",
            "Serviceability",
            f"  Deflection: {status('deflection')} ({values['delta_in']:.3f} in / {values['delta_limit_in']:.3f} in limit, D/C {values['delta_in'] / values['delta_limit_in']:.3f})",
        ]
    )


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Calculate the steel beam 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 a human-readable design summary to stdout")
    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(summary(result) + "\n")
    else:
        output_path.write_text(serialized, encoding="utf-8")
        print(output_path)
    return 0


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