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 .. 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 make_heatmap(aoa: AoA, max_tof: float) -> io.BytesIO:
    logger.info(f"Making heatmap with aoa of {aoa.timestamp}")
    fig = plt.figure()
    ax = fig.add_axes([0, 0, 1, 1], polar=True)
    r = np.linspace(0, max_tof, 100)  # Radius values
    theta = np.linspace(0, np.pi, 50)  # Angle values
    R, Theta = np.meshgrid(r, theta)  # Create a 2D grid of r and theta

    # Compute the function values
    Z = np.log(np.vectorize(aoa.evaluate)(Theta, R))

    ax.pcolormesh(Theta, R, Z, edgecolors="face")
    buf = io.BytesIO()
    fig.savefig(buf, format="jpeg")
    plt.close(fig)

    buf.seek(0)
    return buf


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

    prev_frame = datetime.now()
    while True:
        while (datetime.now() - prev_frame).total_seconds() < 1 / config.HEATMAP_FPS:
            time.sleep(0.01)
        while not aoa_queue.empty():
            logger.debug("Receiving from aoa pipe")
            aoa = aoa_queue.get()
        prev_frame = datetime.now()
        logger.debug(f"Generating heatmap of time {aoa.timestamp}")
        buf = make_heatmap(aoa, max_tof)
        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:
    max_tof_str = request.args.get("max_tof")
    try:
        max_tof = float(max_tof_str)
    except Exception:
        max_tof = 5e-8
    return Response(
        gather_aoa(max_tof), mimetype="multipart/x-mixed-replace; boundary=frame"
    )


def start(conn: "mp.Queue[AoA]") -> 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")