from datetime import datetime
import numpy as np
import io
import base64
from fastapi import APIRouter
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field
from sqlmodel import Session, func, select
from sqlmodel.sql.expression import SelectOfScalar
from db import Chemistry, Player, engine
import networkx as nx

import matplotlib

matplotlib.use("agg")
import matplotlib.pyplot as plt

analysis_router = APIRouter(prefix="/analysis")


C = Chemistry
P = Player


def sociogram_json():
    nodes = []
    necessary_nodes = set()
    links = []
    with Session(engine) as session:
        for p in session.exec(select(P)).fetchall():
            nodes.append({"id": p.name, "appearance": 1})
        subquery = (
            select(C.user, func.max(C.time).label("latest"))
            .where(C.time > datetime(2025, 2, 1, 10))
            .group_by(C.user)
            .subquery()
        )
        statement2 = select(C).join(
            subquery, (C.user == subquery.c.user) & (C.time == subquery.c.latest)
        )
        for c in session.exec(statement2):
            # G.add_node(c.user)
            necessary_nodes.add(c.user)
            for p in c.love:
                # G.add_edge(c.user, p)
                # p_id = session.exec(select(P.id).where(P.name == p)).one()
                necessary_nodes.add(p)
                links.append({"source": c.user, "target": p})
    # nodes = [n for n in nodes if n["name"] in necessary_nodes]
    return JSONResponse({"nodes": nodes, "links": links})


def sociogram_data():
    nodes = []
    links = []
    G = nx.DiGraph()
    with Session(engine) as session:
        for p in session.exec(select(P)).fetchall():
            nodes.append({"id": p.name})
            G.add_node(p.name)
        subquery = (
            select(C.user, func.max(C.time).label("latest"))
            .where(C.time > datetime(2025, 2, 1, 10))
            .group_by(C.user)
            .subquery()
        )
        statement2 = (
            select(C)
            .where(C.user.in_(["Kruse", "Franz", "ck"]))
            .join(subquery, (C.user == subquery.c.user) & (C.time == subquery.c.latest))
        )
        for c in session.exec(statement2):
            for p in c.love:
                G.add_edge(c.user, p)
                links.append({"source": c.user, "target": p})
    return G


class Params(BaseModel):
    node_size: int | None = Field(default=2400, alias="nodeSize")
    font_size: int | None = Field(default=10, alias="fontSize")
    arrow_size: int | None = Field(default=20, alias="arrowSize")
    edge_width: float | None = Field(default=1, alias="edgeWidth")


def sociogram_image(params: Params):
    print(params)
    plt.figure(figsize=(16, 10), facecolor="none")
    ax = plt.gca()
    ax.set_facecolor("none")  # Set the axis face color to none (transparent)
    ax.axis("off")  # Turn off axis ticks and frames

    G = sociogram_data()
    pos = nx.spring_layout(
        G, scale=2, k=1 / np.sqrt(G.number_of_edges()), iterations=50, seed=42
    )
    nx.draw_networkx_nodes(
        G,
        pos,
        node_color="#99ccff",
        edgecolors="#404040",
        linewidths=1,
        node_size=params.node_size,
        alpha=0.86,
    )
    nx.draw_networkx_labels(G, pos, font_size=params.font_size)
    nx.draw_networkx_edges(
        G,
        pos,
        arrows=True,
        edge_color="#404040",
        arrowsize=params.arrow_size,
        node_size=params.node_size,
        width=params.edge_width,
    )

    buf = io.BytesIO()
    plt.savefig(buf, format="png", bbox_inches="tight", dpi=300, transparent=True)
    buf.seek(0)
    encoded_image = base64.b64encode(buf.read()).decode("UTF-8")
    plt.close()

    return {"image": encoded_image}


analysis_router.add_api_route("/json", endpoint=sociogram_json, methods=["GET"])
analysis_router.add_api_route("/image", endpoint=sociogram_image, methods=["POST"])

if __name__ == "__main__":
    with Session(engine) as session:
        statement: SelectOfScalar[P] = select(func.count(P.id))
        print("players in DB: ", session.exec(statement).first())
    G = sociogram_data()
    pos = nx.spring_layout(G, scale=1, k=2, iterations=50, seed=42)
    edges = nx.draw_networkx_edges(
        G,
        pos,
        arrows=True,
        arrowsize=12,
    )
    nx.draw_networkx(
        G, pos, with_labels=True, node_color="#99ccff", font_size=8, node_size=2000
    )
    plt.show()