import React, { useState } from "react";

const ProbabilityTreeWidget: React.FC = () => {
  const [replacement, setReplacement] = useState(false);
  const [initR, setInitR] = useState(3);
  const [initB, setInitB] = useState(4);
  const [hoverPath, setHoverPath] = useState<string | null>(null); // 'RR', 'RB', 'BR', 'BB'

  const total = initR + initB;

  // Level 1 Probs
  const pR = initR / total;
  const pB = initB / total;

  // Level 2 Probs (Given R)
  const r_R = replacement ? initR : Math.max(0, initR - 1);
  const r_Total = replacement ? total : total - 1;
  const pR_R = r_Total > 0 ? r_R / r_Total : 0;
  const pB_R = r_Total > 0 ? 1 - pR_R : 0;

  // Level 2 Probs (Given B)
  const b_B = replacement ? initB : Math.max(0, initB - 1);
  const b_Total = replacement ? total : total - 1;
  const pB_B = b_Total > 0 ? b_B / b_Total : 0;
  const pR_B = b_Total > 0 ? 1 - pB_B : 0;

  // Final Probs
  const pRR = pR * pR_R;
  const pRB = pR * pB_R;
  const pBR = pB * pR_B;
  const pBB = pB * pB_B;

  const fraction = (num: number, den: number) => {
    if (den === 0) return "0";
    return (
      <span className="font-mono bg-white px-1 rounded shadow-sm border border-slate-200 text-xs inline-block mx-1">
        {num}/{den}
      </span>
    );
  };

  const getPathColor = (
    segment:
      | "top"
      | "bottom"
      | "top-top"
      | "top-bottom"
      | "bottom-top"
      | "bottom-bottom",
  ) => {
    const defaultColor = "#cbd5e1"; // Slate 300

    if (!hoverPath) {
      // Default coloring based on branch type
      if (segment.includes("top")) return "#f43f5e"; // Red branches
      if (segment.includes("bottom")) return "#3b82f6"; // Blue branches
      return defaultColor;
    }

    // Highlighting logic based on hoverPath
    if (segment === "top")
      return hoverPath === "RR" || hoverPath === "RB" ? "#f43f5e" : "#f1f5f9";
    if (segment === "bottom")
      return hoverPath === "BR" || hoverPath === "BB" ? "#3b82f6" : "#f1f5f9";

    if (segment === "top-top")
      return hoverPath === "RR" ? "#f43f5e" : "#f1f5f9";
    if (segment === "top-bottom")
      return hoverPath === "RB" ? "#3b82f6" : "#f1f5f9";

    if (segment === "bottom-top")
      return hoverPath === "BR" ? "#f43f5e" : "#f1f5f9";
    if (segment === "bottom-bottom")
      return hoverPath === "BB" ? "#3b82f6" : "#f1f5f9";

    return defaultColor;
  };

  const getStrokeWidth = (segment: string) => {
    if (!hoverPath) return 2;

    if (segment === "top")
      return hoverPath === "RR" || hoverPath === "RB" ? 4 : 1;
    if (segment === "bottom")
      return hoverPath === "BR" || hoverPath === "BB" ? 4 : 1;

    if (segment === "top-top") return hoverPath === "RR" ? 4 : 1;
    if (segment === "top-bottom") return hoverPath === "RB" ? 4 : 1;

    if (segment === "bottom-top") return hoverPath === "BR" ? 4 : 1;
    if (segment === "bottom-bottom") return hoverPath === "BB" ? 4 : 1;

    return 2;
  };

  return (
    <div className="bg-white p-6 rounded-xl shadow-lg border border-slate-200">
      {/* Controls */}
      <div className="flex flex-wrap justify-between items-center mb-6 gap-4">
        <div className="flex gap-4">
          <div className="flex flex-col">
            <label className="text-xs font-bold text-rose-600 uppercase mb-1">
              Red Items
            </label>
            <div className="flex items-center gap-2">
              <button
                onClick={() => setInitR(Math.max(1, initR - 1))}
                className="w-6 h-6 bg-rose-100 text-rose-700 rounded hover:bg-rose-200"
              >
                -
              </button>
              <span className="font-bold w-4 text-center">{initR}</span>
              <button
                onClick={() => setInitR(Math.min(10, initR + 1))}
                className="w-6 h-6 bg-rose-100 text-rose-700 rounded hover:bg-rose-200"
              >
                +
              </button>
            </div>
          </div>
          <div className="flex flex-col">
            <label className="text-xs font-bold text-blue-600 uppercase mb-1">
              Blue Items
            </label>
            <div className="flex items-center gap-2">
              <button
                onClick={() => setInitB(Math.max(1, initB - 1))}
                className="w-6 h-6 bg-blue-100 text-blue-700 rounded hover:bg-blue-200"
              >
                -
              </button>
              <span className="font-bold w-4 text-center">{initB}</span>
              <button
                onClick={() => setInitB(Math.min(10, initB + 1))}
                className="w-6 h-6 bg-blue-100 text-blue-700 rounded hover:bg-blue-200"
              >
                +
              </button>
            </div>
          </div>
        </div>

        <div className="flex bg-slate-100 p-1 rounded-lg">
          <button
            onClick={() => setReplacement(true)}
            className={`px-3 py-1 text-xs font-bold rounded transition-all ${replacement ? "bg-white shadow text-indigo-600" : "text-slate-500"}`}
          >
            With Replacement
          </button>
          <button
            onClick={() => setReplacement(false)}
            className={`px-3 py-1 text-xs font-bold rounded transition-all ${!replacement ? "bg-white shadow text-indigo-600" : "text-slate-500"}`}
          >
            Without Replacement
          </button>
        </div>
      </div>

      <div className="relative h-64 w-full max-w-lg mx-auto select-none">
        <svg width="100%" height="100%" className="overflow-visible">
          {/* Root */}
          <circle cx="20" cy="128" r="6" fill="#64748b" />

          {/* Level 1 Branches */}
          <path
            d="M 20 128 C 50 128, 50 64, 150 64"
            fill="none"
            stroke={getPathColor("top")}
            strokeWidth={getStrokeWidth("top")}
            className="transition-all duration-300"
          />
          <path
            d="M 20 128 C 50 128, 50 192, 150 192"
            fill="none"
            stroke={getPathColor("bottom")}
            strokeWidth={getStrokeWidth("bottom")}
            className="transition-all duration-300"
          />

          {/* Level 1 Labels */}
          <foreignObject x="60" y="70" width="60" height="30">
            <div
              className={`text-center font-bold text-xs ${hoverPath && hoverPath[0] !== "R" ? "text-slate-300" : "text-rose-600"}`}
            >
              {initR}/{total}
            </div>
          </foreignObject>
          <foreignObject x="60" y="150" width="60" height="30">
            <div
              className={`text-center font-bold text-xs ${hoverPath && hoverPath[0] !== "B" ? "text-slate-300" : "text-blue-600"}`}
            >
              {initB}/{total}
            </div>
          </foreignObject>

          {/* Level 1 Nodes */}
          <circle
            cx="150"
            cy="64"
            r="18"
            fill="#f43f5e"
            className={`transition-all ${hoverPath && hoverPath[0] !== "R" ? "opacity-20" : "opacity-100 shadow-md"}`}
          />
          <text
            x="150"
            y="68"
            textAnchor="middle"
            fill="white"
            className={`text-xs font-bold pointer-events-none ${hoverPath && hoverPath[0] !== "R" ? "opacity-20" : ""}`}
          >
            R
          </text>

          <circle
            cx="150"
            cy="192"
            r="18"
            fill="#3b82f6"
            className={`transition-all ${hoverPath && hoverPath[0] !== "B" ? "opacity-20" : "opacity-100 shadow-md"}`}
          />
          <text
            x="150"
            y="196"
            textAnchor="middle"
            fill="white"
            className={`text-xs font-bold pointer-events-none ${hoverPath && hoverPath[0] !== "B" ? "opacity-20" : ""}`}
          >
            B
          </text>

          {/* Level 2 Branches (Top) */}
          <path
            d="M 168 64 L 280 32"
            fill="none"
            stroke={getPathColor("top-top")}
            strokeWidth={getStrokeWidth("top-top")}
            strokeDasharray="4,2"
            className="transition-all duration-300"
          />
          <path
            d="M 168 64 L 280 96"
            fill="none"
            stroke={getPathColor("top-bottom")}
            strokeWidth={getStrokeWidth("top-bottom")}
            strokeDasharray="4,2"
            className="transition-all duration-300"
          />

          {/* Level 2 Top Labels */}
          <foreignObject x="190" y="25" width="60" height="30">
            <div
              className={`text-center font-bold text-xs ${hoverPath === "RR" ? "text-rose-600 scale-110" : "text-slate-400"}`}
            >
              {r_R}/{r_Total}
            </div>
          </foreignObject>
          <foreignObject x="190" y="80" width="60" height="30">
            <div
              className={`text-center font-bold text-xs ${hoverPath === "RB" ? "text-blue-600 scale-110" : "text-slate-400"}`}
            >
              {initB}/{r_Total}
            </div>
          </foreignObject>

          {/* Level 2 Branches (Bottom) */}
          <path
            d="M 168 192 L 280 160"
            fill="none"
            stroke={getPathColor("bottom-top")}
            strokeWidth={getStrokeWidth("bottom-top")}
            strokeDasharray="4,2"
            className="transition-all duration-300"
          />
          <path
            d="M 168 192 L 280 224"
            fill="none"
            stroke={getPathColor("bottom-bottom")}
            strokeWidth={getStrokeWidth("bottom-bottom")}
            strokeDasharray="4,2"
            className="transition-all duration-300"
          />

          {/* Level 2 Bottom Labels */}
          <foreignObject x="190" y="150" width="60" height="30">
            <div
              className={`text-center font-bold text-xs ${hoverPath === "BR" ? "text-rose-600 scale-110" : "text-slate-400"}`}
            >
              {initR}/{b_Total}
            </div>
          </foreignObject>
          <foreignObject x="190" y="210" width="60" height="30">
            <div
              className={`text-center font-bold text-xs ${hoverPath === "BB" ? "text-blue-600 scale-110" : "text-slate-400"}`}
            >
              {b_B}/{b_Total}
            </div>
          </foreignObject>

          {/* Outcomes (Interactive Targets) */}
          <g
            className="cursor-pointer"
            onMouseEnter={() => setHoverPath("RR")}
            onMouseLeave={() => setHoverPath(null)}
          >
            <text
              x="300"
              y="36"
              className={`text-xs font-bold transition-all ${hoverPath === "RR" ? "fill-rose-600 text-base" : "fill-slate-500"}`}
            >
              RR: {(pRR * 100).toFixed(1)}%
            </text>
            <rect x="290" y="20" width="80" height="20" fill="transparent" />
          </g>

          <g
            className="cursor-pointer"
            onMouseEnter={() => setHoverPath("RB")}
            onMouseLeave={() => setHoverPath(null)}
          >
            <text
              x="300"
              y="100"
              className={`text-xs font-bold transition-all ${hoverPath === "RB" ? "fill-indigo-600 text-base" : "fill-slate-500"}`}
            >
              RB: {(pRB * 100).toFixed(1)}%
            </text>
            <rect x="290" y="85" width="80" height="20" fill="transparent" />
          </g>

          <g
            className="cursor-pointer"
            onMouseEnter={() => setHoverPath("BR")}
            onMouseLeave={() => setHoverPath(null)}
          >
            <text
              x="300"
              y="164"
              className={`text-xs font-bold transition-all ${hoverPath === "BR" ? "fill-indigo-600 text-base" : "fill-slate-500"}`}
            >
              BR: {(pBR * 100).toFixed(1)}%
            </text>
            <rect x="290" y="150" width="80" height="20" fill="transparent" />
          </g>

          <g
            className="cursor-pointer"
            onMouseEnter={() => setHoverPath("BB")}
            onMouseLeave={() => setHoverPath(null)}
          >
            <text
              x="300"
              y="228"
              className={`text-xs font-bold transition-all ${hoverPath === "BB" ? "fill-blue-600 text-base" : "fill-slate-500"}`}
            >
              BB: {(pBB * 100).toFixed(1)}%
            </text>
            <rect x="290" y="215" width="80" height="20" fill="transparent" />
          </g>
        </svg>
      </div>

      {/* Calculation Panel */}
      <div
        className={`p-4 rounded-lg border text-sm mt-4 transition-colors ${hoverPath ? "bg-amber-50 border-amber-200 text-amber-900" : "bg-slate-50 border-slate-100 text-slate-400"}`}
      >
        {!hoverPath ? (
          <p className="text-center italic">
            Hover over an outcome (e.g., RR) to see the calculation.
          </p>
        ) : (
          <>
            <p className="font-bold mb-1">
              Calculation for{" "}
              <span className="font-mono bg-white px-1 rounded border border-amber-200">
                {hoverPath}
              </span>
              ({hoverPath[0] === "R" ? "Red" : "Blue"} then{" "}
              {hoverPath[1] === "R" ? "Red" : "Blue"}):
            </p>
            <div className="flex flex-wrap items-center gap-2 font-mono text-lg mt-2 justify-center sm:justify-start">
              {/* First Draw */}
              <span>P({hoverPath[0]})</span>
              <span>×</span>
              <span>
                P({hoverPath[1]} | {hoverPath[0]})
              </span>
              <span>=</span>

              {/* Numbers */}
              {fraction(hoverPath[0] === "R" ? initR : initB, total)}
              <span>×</span>
              {fraction(
                hoverPath === "RR"
                  ? r_R
                  : hoverPath === "RB"
                    ? initB
                    : hoverPath === "BR"
                      ? initR
                      : b_B,
                hoverPath[0] === "R" ? r_Total : b_Total,
              )}
              <span>=</span>

              {/* Result */}
              <strong className="text-amber-700">
                {fraction(
                  (hoverPath[0] === "R" ? initR : initB) *
                    (hoverPath === "RR"
                      ? r_R
                      : hoverPath === "RB"
                        ? initB
                        : hoverPath === "BR"
                          ? initR
                          : b_B),
                  total * (hoverPath[0] === "R" ? r_Total : b_Total),
                )}
              </strong>
            </div>
            {!replacement && hoverPath[0] === hoverPath[1] && (
              <p className="text-xs mt-3 text-rose-600 font-bold bg-white p-2 rounded inline-block border border-rose-100">
                ⚠ Notice: The numerator decreased because we kept the first{" "}
                {hoverPath[0] === "R" ? "Red" : "Blue"} item!
              </p>
            )}
          </>
        )}
      </div>
    </div>
  );
};

export default ProbabilityTreeWidget;