import io
import logging
import multiprocessing as mp
import time
from datetime import datetime
from typing import Generator

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import numpy.typing as npt
from flask import Flask, Response, render_template, request
from flask_sock import Sock
from simple_websocket import Server

from ..config import config
from ..processing.aoa import AoA

matplotlib.use("agg")

app = Flask(__name__)
sock = Sock(app)
aoa_queue = None

logger = logging.getLogger(__name__)


data: npt.NDArray[np.complex128] = np.array([], dtype=complex)
aoa: AoA = AoA()


@app.route("/preprocessed")
def preprocessed() -> str:
    return render_template("preprocessed.html")


Subscriber = Server
subscriber_settings: dict[Subscriber, tuple[int, int, int]] = {}


@sock.route("/data")
def get_data(sock: Subscriber) -> None:
    while True:
        msg = sock.receive()
        if len(msg.split()) != 3:
            break
        subcarrier, rx, tx = map(int, msg.split())
        subscriber_settings[sock] = (subcarrier, rx, tx)


def add_data(
    raw_data: npt.NDArray[np.complex128], new_data: npt.NDArray[np.complex128]
) -> None:
    if config.VISUALISE_RAW:
        magn = np.abs(raw_data)
        phase = np.angle(raw_data)

        new_mag = np.abs(new_data)
        new_phase = np.angle(new_data)

        fig, axs = plt.subplots(2, 2)
        axs[0, 0].plot(magn[:, 0, 0], c="b")
        axs[0, 0].plot(magn[:, 1, 0], c="orange")
        axs[1, 0].plot(new_mag[:, 0, 0], c="b")
        axs[1, 0].plot(new_mag[:, 1, 0], c="orange")
        axs[0, 1].plot(phase[:, 0, 0], c="b")
        axs[0, 1].plot(phase[:, 1, 0], c="orange")
        axs[1, 1].plot(new_phase[:, 0, 0], c="b")
        axs[1, 1].plot(new_phase[:, 1, 0], c="orange")
        fig.savefig("/tmp/plot.png")
        plt.close(fig)

    global data
    if data.size == 0:
        data = np.expand_dims(new_data, axis=0)
    else:
        data = np.concat([data, np.expand_dims(new_data, axis=0)], axis=0)

        # Only keep latest 100 entries
        if data.shape[0] > 100:
            data = data[-100:]
    to_remove: list[Subscriber] = []
    for subscriber in subscriber_settings:
        try:
            subcarrier, rx, tx = subscriber_settings[subscriber]
            subscriber.send(new_data.real[subcarrier, rx, tx])
        except Exception as e:
            to_remove.append(subscriber)
            print(e)

    for subscriber in to_remove:
        del subscriber_settings[subscriber]


def plot_heatmap(heatmap: npt.NDArray[np.float32]) -> io.BytesIO:
    logger.info(f"Making heatmap with aoa of {aoa.timestamp}")
    fig = plt.figure()
    ax = fig.add_axes([0.1, 0.1, 0.9, 0.9])  # , polar=True)
    r = np.linspace(
        0, config.music.heatmap.tof_max, config.music.heatmap.tof_resolution
    )
    theta = np.linspace(0, np.pi, config.music.heatmap.theta_resolution)  # Angle values

    mesh = ax.pcolormesh(theta, r, heatmap, edgecolors="face", vmin=0, vmax=50)

    fig.colorbar(mesh, ax=ax)
    buf = io.BytesIO()
    fig.savefig(buf, format="jpeg")
    plt.close(fig)

    buf.seek(0)
    return buf


def gather_aoa() -> Generator[bytes, None, None]:
    assert aoa_queue is not None

    while True:
        logger.debug("Receiving from aoa pipe")
        heatmap = aoa_queue.get()
        buf = plot_heatmap(heatmap)
        yield (b"--frame\r\nContent-Type: image/jpeg\r\n\r\n" + buf.read() + b"\r\n")
        buf.close()


@app.route("/aoa_tof")
def aoa_tof() -> Response:
    return Response(gather_aoa(), mimetype="multipart/x-mixed-replace; boundary=frame")


def start(conn: "mp.Queue[npt.NDArray[np.float32]]") -> None:
    global app, aoa_queue

    aoa_queue = conn
    app.run(debug=True, use_reloader=False, host="0.0.0.0")
    logger.info("Visualisation server shut down")