from typing import TYPE_CHECKING, Any, Callable, List, Optional, Protocol

import json
import math


if TYPE_CHECKING:

    class _PyScript(Protocol):
        document: Any

    pyscript: _PyScript
else:
    import pyscript

from solver_core import MultisetGraph


KG_UNITS_PER_KG = 4
LB_PER_KG = 2.2046226218
BAR_KG = 20.0
BAR_UNITS = int(BAR_KG * KG_UNITS_PER_KG)


def lbs_to_kg_units(lbs: float) -> int:
    # total weight, in quarter-kg units, rounded up
    kg = lbs / LB_PER_KG
    return int(math.ceil(kg * KG_UNITS_PER_KG))


def reachable_table(nums: List[int]) -> List[bool]:
    total = sum(nums)
    reachable = [False] * (total + 1)
    reachable[0] = True
    for x in nums:
        for t in range(total, x - 1, -1):
            if reachable[t - x]:
                reachable[t] = True
    return reachable


def ceil_reachable(reachable: List[bool], target: int) -> int:
    if target <= 0:
        return 0
    if target >= len(reachable):
        return len(reachable) - 1
    for t in range(target, len(reachable)):
        if reachable[t]:
            return t
    return len(reachable) - 1


def format_kg_plate_side(pair_units: int) -> str:
    q = pair_units // 2  # quarter-kg units per side
    whole = q // 4
    rem = q % 4

    if rem == 0:
        s = str(whole)
        # interpunct-pad 1-digit integers (mainly "5")
        if len(s) == 1:
            s = s + "·"
        return s
    elif rem == 1:
        return f"{whole}¼" if whole else "¼"
    elif rem == 2:
        return f"{whole}½" if whole else "½"
    else:  # rem == 3
        return f"{whole}¾" if whole else "¾"


def format_kg_total(total_units: int) -> str:
    kg = total_units / 4.0
    return f"{kg:5.1f} kg"  # e.g. " 92.5 kg", "115.0 kg"


def doit(
    plates,
    targets,
    reps,
    *,
    bar_units: int = 45,
    format_plate_side: Optional[Callable[[int], str]] = None,
    format_total: Optional[Callable[[int], str]] = None,
    hints: Optional[List[str]] = None,
    unit_mode: str = "lb",
):
    graph = MultisetGraph(plates, targets)
    graph.build_layers()
    _, shortest_path = graph.dijkstra()
    lines = []

    # Default lb formatting (matches your current behavior)
    def default_format_plate_side(n: int) -> str:
        if n == 5:
            return "2½"
        elif n == 10:
            return "5·"
        else:
            return str(n // 2)

    def default_format_total(total_units: int) -> str:
        return str(total_units)

    if format_plate_side is None:
        format_plate_side = default_format_plate_side
    if format_total is None:
        format_total = default_format_total

    lines.append("")
    for prev_node, node in zip(shortest_path, shortest_path[1:-1]):
        if node.layer == 0:
            continue

        # Find the longest common prefix
        list_a = prev_node.elements
        list_b = node.elements
        p = 0
        while p < len(list_a) and p < len(list_b) and list_a[p] == list_b[p]:
            p += 1

        pops = list_a[p:][::-1]
        pushes = list_b[p:]

        set_lines = []

        # Keep existing coloring behavior for lb mode.
        # For kg mode later, you can either accept string plates (no color) or tweak this.
        def _parse_plate_value(s: str) -> Optional[float]:
            if s.endswith("·"):
                try:
                    return float(s[:-1])
                except ValueError:
                    pass

            # "2½" is ambiguous: 2.5 lb (small-green) vs 2.5 kg (amber bucket)
            if s == "2½":
                return 2.5

            # kg fraction glyphs like "1¼", "12½", etc.
            frac_map = {"¼": 0.25, "½": 0.5, "¾": 0.75}
            if len(s) >= 1 and s[-1] in frac_map:
                try:
                    whole = int(s[:-1]) if s[:-1] != "" else 0
                    return whole + frac_map[s[-1]]
                except ValueError:
                    pass

            try:
                return float(s)
            except ValueError:
                return None

        def color_plate(plate_str: str) -> str:
            v = _parse_plate_value(plate_str)
            if v is None:
                return plate_str

            # Mode-specific color maps (per-side values)
            LB_COLORS = {
                2.5: "#2dc84d",
                5.0: "#ffb703",
                10.0: "#00b4d8",
                25.0: "#e63946",
                35.0: "#3fae2a",
                45.0: "#0077b6",
            }

            # Map kg plates to "closest lb family" colors:
            # 1.25kg ≈ 2.5lb (green)
            # 2.5kg  ≈ 5lb   (amber)
            # 5kg    ≈ 10lb  (cyan)
            # 10kg   ≈ 25lb  (red)   (closer than 35lb)
            # 15kg   ≈ 35lb  (green2)
            # 20kg   ≈ 45lb  (blue)
            # 25kg   ≈ 55lb  (no direct lb class; pick blue or red—I'll pick blue-ish as "big")
            KG_COLORS = {
                1.25: "#2dc84d",
                2.5: "#ffb703",
                5.0: "#00b4d8",
                10.0: "#e63946",
                15.0: "#3fae2a",
                20.0: "#0077b6",
                25.0: "#0077b6",
            }

            # Choose map
            cmap = KG_COLORS if unit_mode == "kg" else LB_COLORS

            # Float safety: compare with tolerance
            for key, color in cmap.items():
                if abs(v - key) < 1e-9:
                    return f"<span style='color: {color}'>{plate_str}</span>"

            return plate_str

        plate_side_strs = [format_plate_side(x) for x in node.elements]
        plate_list = "<span style='color: #888888'>–</span>".join(
            [color_plate(s) for s in plate_side_strs]
        )

        set_lines.append("<ul>")
        if node.layer == 1:
            pass
        else:
            for pop in pops:
                set_lines.append(
                    f"<li style='color: #ee2200'> pop {format_plate_side(pop)}</li>"
                )
            for push in pushes:
                set_lines.append(
                    f"<li style='color: #0000bb'> push {format_plate_side(push)}</li>"
                )
        set_lines.append("</ul>")

        total_units = sum(node.elements) + bar_units
        hint = ""
        if hints is not None and 0 <= (node.layer - 1) < len(hints):
            hint = f" {hints[node.layer - 1]}"

        set_lines.append(
            "<h2><span style='font-family: Menlo, monospace'>"
            f"#{node.layer}: {format_total(total_units)} ({plate_list}){hint}"
            "</span></h2>"
        )

        def rep_color(num_reps):
            try:
                return [
                    "#f0f921",
                    "#fada24",
                    "#febd2a",
                    "#fba238",
                    "#f48849",
                    "#e97158",
                    "#db5c68",
                    "#cc4778",
                ][num_reps]
            except IndexError:
                return "#cc4778"

        num_reps = reps[node.layer - 1]
        set_lines.append(
            f"<span style='background-color: {rep_color(num_reps)}'>"
            f"{num_reps} rep{'s' if num_reps != 1 else ''}"
            "</span>"
        )
        lines.append("".join(set_lines))

    return "".join(lines)


def set_output(contents):
    output_div = pyscript.document.querySelector("#output")
    output_div.innerHTML = contents


def calculate_weights(event):
    input_text = pyscript.document.querySelector("#weights")
    set_output("Loading...")

    mode = pyscript.document.querySelector('input[name="unit-mode"]:checked').value

    if mode == "kg":
        if input_text.value == "":
            weights_total_lbs = [200, 225, 250, 240, 225, 210, 200, 185, 170]
            reps = [1, 2, 3, 4, 5, 6, 7, 8, 9]
        else:
            weights_total_lbs = [
                int(weight)
                for weight_and_reps in input_text.value.split(",")
                for weight, reps in [weight_and_reps.split("*")]
            ]
            reps = [
                int(reps)
                for weight_and_reps in input_text.value.split(",")
                for weight, reps in [weight_and_reps.split("*")]
            ]

        plate_list_kg = pyscript.document.querySelector("#plate-list-kg")
        plates_kg_units = json.loads(plate_list_kg.value)

        reachable = reachable_table(plates_kg_units)

        adjusted_targets = []
        hints = []

        for lbs in weights_total_lbs:
            target_total_units = lbs_to_kg_units(lbs)
            target_plate_units = max(0, target_total_units - BAR_UNITS)
            adjusted_plate_units = ceil_reachable(reachable, target_plate_units)
            adjusted_targets.append(adjusted_plate_units)

            achieved_total_kg = (BAR_UNITS + adjusted_plate_units) / KG_UNITS_PER_KG
            achieved_lb = achieved_total_kg * LB_PER_KG

            hints.append(f"({lbs} → {achieved_lb:.1f} lb)")

        output = doit(
            plates_kg_units,
            adjusted_targets,
            reps,
            bar_units=BAR_UNITS,
            format_plate_side=format_kg_plate_side,
            format_total=format_kg_total,
            hints=hints,
            unit_mode="kg",
        )
        set_output(output)
        return

    # lb mode: existing behavior (targets are plate-only lbs, i.e. total-45)
    if input_text.value == "":
        weights = [
            weight - 45 for weight in [200, 225, 250, 240, 225, 210, 200, 185, 170]
        ]
        reps = [1, 2, 3, 4, 5, 6, 7, 8, 9]
    else:
        weights = [
            int(weight) - 45
            for weight_and_reps in input_text.value.split(",")
            for weight, reps in [weight_and_reps.split("*")]
        ]
        reps = [
            int(reps)
            for weight_and_reps in input_text.value.split(",")
            for weight, reps in [weight_and_reps.split("*")]
        ]

    plate_list = pyscript.document.querySelector("#plate-list")
    plates = json.loads(plate_list.value)

    output = doit(plates, weights, reps)
    set_output(output)