"""Utilities for canonizing and saving graphs."""
import base64
import tempfile
from collections import defaultdict
from pathlib import Path
from typing import Any, cast, TypedDict

from sage.all import Graph, save, load  # type: ignore[attr-defined]  # pylint: disable=no-name-in-module
from lib.tutte_embedding import tutte_embedding
from lib.planar_embedding import outer_face

DIR = Path(__file__).parent.parent
PALETTE = ['red', 'blue', 'green', 'yellow']

VertexColoring = dict[Any, Any]


class ColoredGraphId(TypedDict):
    """Canonical id representing a colored graph"""
    graph_id: str
    coloring_id: str


def canonize_colored_graph(g: Graph, coloring: VertexColoring) -> tuple[Graph, VertexColoring, ColoredGraphId]:
    """Return a new canonical graph, new canonical coloring, and a ColoredGraphId without mutating inputs"""
    canonical, cert = cast(
        tuple[Graph, dict[Any, int]],
        g.canonical_label(certificate=True),
    )
    graph_id = base64.urlsafe_b64encode(
        canonical.graph6_string().encode()
    ).decode()

    color_seq = [0] * g.order()
    for orig_v, canon_idx in cert.items():
        color_seq[canon_idx] = coloring[orig_v]

    canonical_coloring = {canon_idx: color for canon_idx, color in enumerate(color_seq)}
    coloring_id = base64.urlsafe_b64encode(bytes(color_seq)).decode()
    return canonical, canonical_coloring, ColoredGraphId(graph_id=graph_id, coloring_id=coloring_id)


def plot_colored_graph_to_data_uri(g: Graph, coloring: VertexColoring) -> str:
    vertex_colors: defaultdict[str, list[Any]] = defaultdict(list)
    for v, c in coloring.items():
        vertex_colors[PALETTE[c]].append(v)
    if g.get_pos() is None:
        g.is_planar(set_embedding=True, set_pos=True)
    g.set_pos(tutte_embedding(g, outer_face(g)))
    with tempfile.NamedTemporaryFile(suffix='.png', delete=True) as f:
        g.plot(vertex_colors=dict(vertex_colors)).save(f.name)
        png_bytes = Path(f.name).read_bytes()
    return f"data:image/png;base64,{base64.b64encode(png_bytes).decode()}"


def canonize_and_save_graph(g: Graph) -> tuple[Graph, Path]:
    """Canonize g, save to data/graphs/<graph_id>/graph.sobj, and return the canonical graph and its folder."""
    canonical = cast(Graph, g.canonical_label())
    graph_id = base64.urlsafe_b64encode(canonical.graph6_string().encode()).decode()
    out_dir = DIR / "data" / "graphs" / graph_id
    out_dir.mkdir(parents=True, exist_ok=True)
    print(graph_id)
    canonical.is_planar(set_embedding=True, set_pos=True)
    canonical.plot().save(out_dir / 'graph.png')
    save(canonical, str(out_dir / 'graph'))
    return canonical, out_dir


def save_colored_graph(g: Graph, coloring: VertexColoring) -> tuple[Graph, VertexColoring, ColoredGraphId]:
    """
    Canonize g and coloring, save to disk, and return the canonical forms with id.
    If already saved, load and return the cached graph.
    """
    g_canon, canonical_coloring, cid = canonize_colored_graph(g, coloring)
    out_dir = DIR / "data" / "graphs" / cid['graph_id'] / cid['coloring_id']
    if (out_dir / "graph.sobj").exists():
        g_canon = cast(Graph, load(str(out_dir / 'graph')))
        return g_canon, canonical_coloring, cid
    g_canon.is_planar(set_embedding=True, set_pos=True)
    vertex_colors: defaultdict[str, list[Any]] = defaultdict(list)
    for v, c in canonical_coloring.items():
        vertex_colors[PALETTE[c]].append(v)
    out_dir.mkdir(parents=True, exist_ok=True)
    g_canon.plot(
        vertex_colors=dict(vertex_colors),
        title=f"graph: {cid['graph_id']} coloring: {cid['coloring_id']}",
    ).save(out_dir / 'graph.png')
    save(g_canon, str(out_dir / 'graph'))
    return g_canon, canonical_coloring, cid
