from __future__ import annotations

import json
from html import escape
from pathlib import Path
from typing import Any

import plotly.graph_objects as go
import plotly.io as pio
import torch


def _build_hover_text(
    item_id: str,
    metadata: dict[str, Any] | None,
    hover_fields: list[str] | None,
) -> str:
    rows = [f"id: {item_id}"]
    if metadata:
        if hover_fields:
            for key in hover_fields:
                if key in metadata:
                    rows.append(f"{key}: {metadata[key]}")
        else:
            for key, value in metadata.items():
                if key == "id":
                    continue
                rows.append(f"{key}: {value}")
                if len(rows) >= 7:
                    break
    return "<br>".join(escape(str(x)) for x in rows)


def create_3d_figure(
    coords_3d: torch.Tensor,
    ids: list[str],
    metadata_by_id: dict[str, dict[str, Any]],
    image_paths_by_id: dict[str, str],
    explained_variance_ratio: list[float],
    title: str,
    hover_fields: list[str] | None,
    point_size: int,
    color_by: str | None,
    legend_title: str | None,
) -> tuple[go.Figure, list[dict[str, Any]]]:
    if coords_3d.ndim != 2 or coords_3d.shape[1] != 3:
        raise ValueError(f"Expected PCA coordinates shape (N, 3), got {tuple(coords_3d.shape)}")

    coords_np = coords_3d.detach().cpu().numpy()

    hover_text = []
    payloads: list[dict[str, Any]] = []
    groups: list[str] = []
    for idx, item_id in enumerate(ids):
        metadata = metadata_by_id.get(item_id)
        image_path = image_paths_by_id.get(item_id)
        if color_by:
            group_value = (metadata or {}).get(color_by, "<missing>")
            group = str(group_value)
        else:
            group = "All points"

        hover_text.append(
            _build_hover_text(item_id=item_id, metadata=metadata, hover_fields=hover_fields)
        )
        groups.append(group)
        payloads.append(
            {
                "id": item_id,
                "index": idx,
                "group": group,
                "metadata": metadata or {},
                "image_path": image_path,
            }
        )

    ratio_pct = [x * 100.0 for x in explained_variance_ratio]
    axis_titles = [
        f"PC1 ({ratio_pct[0]:.2f}%)" if len(ratio_pct) > 0 else "PC1",
        f"PC2 ({ratio_pct[1]:.2f}%)" if len(ratio_pct) > 1 else "PC2",
        f"PC3 ({ratio_pct[2]:.2f}%)" if len(ratio_pct) > 2 else "PC3",
    ]

    grouped_indices: dict[str, list[int]] = {}
    for idx, group in enumerate(groups):
        if group not in grouped_indices:
            grouped_indices[group] = []
        grouped_indices[group].append(idx)

    traces: list[go.Scatter3d] = []
    for group_name, idxs in grouped_indices.items():
        traces.append(
            go.Scatter3d(
                x=coords_np[idxs, 0],
                y=coords_np[idxs, 1],
                z=coords_np[idxs, 2],
                mode="markers",
                text=[hover_text[i] for i in idxs],
                customdata=idxs,
                hovertemplate="%{text}<extra></extra>",
                name=group_name,
                legendgroup=group_name,
                showlegend=bool(color_by),
                marker={"size": point_size, "opacity": 0.86},
            )
        )

    fig = go.Figure(data=traces)

    fig.update_layout(
        title=title,
        template="plotly_white",
        margin={"l": 8, "r": 8, "b": 8, "t": 42},
        scene={
            "xaxis": {"title": axis_titles[0]},
            "yaxis": {"title": axis_titles[1]},
            "zaxis": {"title": axis_titles[2]},
        },
        legend={
            "title": {"text": legend_title or color_by or "Group"},
        },
        height=760,
    )

    return fig, payloads


def save_interactive_html(
    fig: go.Figure,
    payloads: list[dict[str, Any]],
    output_html: str | Path,
) -> None:
    output_html = Path(output_html)

    plot_div = pio.to_html(
        fig,
        include_plotlyjs="cdn",
        full_html=False,
        div_id="citall_plot",
    )

    payloads_json = json.dumps(payloads)

    html = f"""<!doctype html>
<html lang=\"en\">
<head>
  <meta charset=\"utf-8\" />
  <meta name=\"viewport\" content=\"width=device-width, initial-scale=1\" />
  <title>citall explorer</title>
  <style>
    :root {{
      --bg: #f4f6f7;
      --panel: #ffffff;
      --ink: #1d2433;
      --muted: #5f697d;
      --accent: #0d8a74;
      --line: #d7dde5;
    }}
    * {{ box-sizing: border-box; }}
    body {{
      margin: 0;
      background: radial-gradient(circle at 10% 10%, #ffffff 0%, var(--bg) 50%, #e9eef4 100%);
      color: var(--ink);
      font: 14px/1.45 "Avenir Next", "Segoe UI", sans-serif;
    }}
    .wrap {{
      display: grid;
      grid-template-columns: minmax(0, 2fr) minmax(280px, 1fr);
      gap: 16px;
      padding: 16px;
      min-height: 100vh;
    }}
    .plot-panel, .info-panel {{
      background: var(--panel);
      border: 1px solid var(--line);
      border-radius: 16px;
      box-shadow: 0 10px 28px rgba(27, 43, 57, 0.08);
      overflow: hidden;
    }}
    .plot-panel {{ padding: 6px; }}
    .info-panel {{
      padding: 14px;
      display: grid;
      align-content: start;
      gap: 10px;
    }}
    .label {{
      color: var(--muted);
      text-transform: uppercase;
      font-size: 11px;
      letter-spacing: 0.08em;
      margin: 0;
    }}
    .value {{ margin: 0; word-break: break-word; }}
    .meta {{
      display: grid;
      gap: 8px;
      max-height: 40vh;
      overflow: auto;
      border-top: 1px solid var(--line);
      padding-top: 10px;
    }}
    .meta-row {{
      border-bottom: 1px dashed var(--line);
      padding-bottom: 6px;
    }}
    .preview {{
      width: 100%;
      max-height: 260px;
      object-fit: contain;
      border: 1px solid var(--line);
      border-radius: 10px;
      background: #f7f9fb;
      display: none;
    }}
    @media (max-width: 1024px) {{
      .wrap {{ grid-template-columns: 1fr; }}
      .info-panel {{ order: -1; }}
    }}
  </style>
</head>
<body>
  <div class=\"wrap\">
    <section class=\"plot-panel\">{plot_div}</section>
    <aside class=\"info-panel\">
      <p class=\"label\">Selected point</p>
      <h2 id=\"point-id\" class=\"value\">none</h2>
      <p id=\"point-index\" class=\"value\">Click a point to inspect metadata.</p>
      <img id=\"point-image\" class=\"preview\" alt=\"preview\" />
      <div id=\"point-meta\" class=\"meta\"></div>
    </aside>
  </div>

  <script>
    const payloads = {payloads_json};
    const plot = document.getElementById('citall_plot');
    const pointId = document.getElementById('point-id');
    const pointIndex = document.getElementById('point-index');
    const pointImage = document.getElementById('point-image');
    const pointMeta = document.getElementById('point-meta');

    function clearMeta() {{
      while (pointMeta.firstChild) pointMeta.removeChild(pointMeta.firstChild);
    }}

    function renderMeta(metadata) {{
      clearMeta();
      const entries = Object.entries(metadata || {{}});
      if (!entries.length) {{
        const empty = document.createElement('p');
        empty.className = 'value';
        empty.textContent = 'No metadata available for this point.';
        pointMeta.appendChild(empty);
        return;
      }}
      for (const [key, value] of entries) {{
        const row = document.createElement('div');
        row.className = 'meta-row';

        const label = document.createElement('p');
        label.className = 'label';
        label.textContent = key;

        const val = document.createElement('p');
        val.className = 'value';
        val.textContent = typeof value === 'string' ? value : JSON.stringify(value);

        row.appendChild(label);
        row.appendChild(val);
        pointMeta.appendChild(row);
      }}
    }}

    function renderPoint(pointIndexValue) {{
      const payload = payloads[pointIndexValue];
      if (!payload) return;

      pointId.textContent = payload.id ?? 'unknown';
      pointIndex.textContent = `index: ${{payload.index}}`;
      renderMeta(payload.metadata || {{}});

      if (payload.image_path) {{
        pointImage.src = payload.image_path;
        pointImage.style.display = 'block';
      }} else {{
        pointImage.removeAttribute('src');
        pointImage.style.display = 'none';
      }}
    }}

    if (plot && plot.on) {{
      plot.on('plotly_click', function(eventData) {{
        if (!eventData || !eventData.points || !eventData.points.length) return;
        const point = eventData.points[0];
        let idx = point.customdata;
        if (Array.isArray(idx)) idx = idx[0];
        if (typeof idx !== 'number') idx = point.pointIndex;
        renderPoint(idx);
      }});
    }}
  </script>
</body>
</html>
"""

    output_html.parent.mkdir(parents=True, exist_ok=True)
    output_html.write_text(html, encoding="utf-8")
