from pathlib import Path
from typing import cast

import grpc
import matplotlib.pyplot as plt
import numpy as np
import typer

from ..generated import figure_pb2, figure_pb2_grpc

app = typer.Typer()


@app.command()
def list() -> None:
    """
    List all figures from the gRPC service.
    """
    print("Connecting to the server...")
    with grpc.insecure_channel("localhost:50051") as channel:
        stub = figure_pb2_grpc.FigureServiceStub(channel)
        print("Retrieving available figures...\n")
        figures = [x for x in stub.GetFigure(figure_pb2.FigureRequest())]

        if not figures:
            print("No figures found.")
        else:
            for i, figure in enumerate(figures, 1):
                print(f"{i}. Figure: {figure.title} (ID: {figure.uuid})")

            choice = typer.prompt("Which figure to extract?", type=int)
            assert isinstance(choice, int)
            if 1 <= choice <= len(figures):
                selected_figure = figures[choice - 1]
                print(f"You selected: {selected_figure.title}")
                extract(selected_figure.uuid)
            else:
                print("Invalid selection. Exiting.")


def get_figure(uuid: str) -> tuple[figure_pb2.Figure, figure_pb2.FigureData]:
    """
    Get the figure data for a specific UUID.
    """
    print(f"Connecting to the server to extract data for UUID: {uuid}...")
    with grpc.insecure_channel("localhost:50051") as channel:
        stub = figure_pb2_grpc.FigureServiceStub(channel)
        all_figures = {x.uuid: x for x in stub.GetFigure(figure_pb2.FigureRequest())}
        figure_data = stub.GetFigureUpdate(figure_pb2.FigureDataRequest(uuid=uuid))
        for data in figure_data:
            return (all_figures[uuid], data)
    raise ValueError(f"Figure with UUID {uuid} not found.")


@app.command()
def extract(uuid: str, output: Path | None = None) -> None:
    """
    Extract data for a specific figure identified by its UUID.
    """
    _, figure = get_figure(uuid)
    if not output:
        output = cast(Path, typer.prompt("Enter output file path:", type=Path))
    with open(output, "wb") as f:
        f.write(figure.SerializeToString())


@app.command()
def plot(uuid: str) -> None:
    """
    Plot the figure data for a specific UUID.
    """
    figure, data = get_figure(uuid)
    if data.line:
        for line in data.line.lines:
            plt.plot(
                line.x if line.x else np.arange(len(line.y)), line.y, label=line.label
            )
        plt.xlabel(figure.x_label)
        plt.ylabel(figure.y_label)
        plt.title(figure.title)
    elif figure.heatmap:
        plt.imshow(figure.heatmap.data, cmap="hot", interpolation="nearest")
    elif figure.histogram:
        plt.hist(figure.histogram.data, bins=figure.histogram.bins)
    plt.show()


if __name__ == "__main__":
    app()