import psycopg
from psycopg.rows import dict_row, class_row
from domain import Centroid, Song


def choose_centroids(conn: psycopg.Connection, k: int = 1000) -> dict[int, Centroid]:
    final_dict: dict[int, Centroid] = dict()
    QUERY = """
select 
	row_number() over () as id, *
	from (	
		select 
			"Tempo",
			"Zcr",
			"MeanSpectralCentroid",
			"StdDevSpectralCentroid",
			"MeanSpectralRolloff",
			"StdDevSpectralRolloff",
			"MeanSpectralFlatness",
			"StdDevSpectralFlatness",
			"MeanLoudness",
			"StdDevLoudness",
			"Chroma1",
			"Chroma2",
			"Chroma3",
			"Chroma4",
			"Chroma5",
			"Chroma6",
			"Chroma7",
			"Chroma8",
			"Chroma9",
			"Chroma10"
		FROM bliss_tracks bt
		ORDER BY RANDOM() 
		LIMIT %(k)s
) with_ids;"""
    with conn.cursor(row_factory=class_row(Centroid)) as cur:
        cur.execute(QUERY, {"k": k})
        for centroid in cur.fetchall():
            final_dict[centroid.id] = centroid

        return final_dict


def get_all_songs(conn: psycopg.Connection) -> list[Song]:
    QUERY = """
    select 
        "File",
        "Tempo",
        "Zcr",
        "MeanSpectralCentroid",
        "StdDevSpectralCentroid",
        "MeanSpectralRolloff",
        "StdDevSpectralRolloff",
        "MeanSpectralFlatness",
        "StdDevSpectralFlatness",
        "MeanLoudness",
        "StdDevLoudness",
        "Chroma1",
        "Chroma2",
        "Chroma3",
        "Chroma4",
        "Chroma5",
        "Chroma6",
        "Chroma7",
        "Chroma8",
        "Chroma9",
        "Chroma10"
    FROM bliss_tracks;"""
    with conn.cursor(row_factory=class_row(Song)) as cur:
        cur.execute(QUERY)
        return cur.fetchall()


def get_conn() -> psycopg.Connection:
    from os import environ

    return psycopg.connect(
        f"dbname=postgres user=postgres port=5555 host=192.168.1.68 password={environ['PGPASSWORD']}"
    )