"""Functions for accessing the database relating to datasets."""
import MySQLdb as mdb
from MySQLdb.cursors import DictCursor

def geno_datasets_by_species_and_population(
        conn: mdb.Connection,
        speciesid: int,
        populationid: int) -> tuple[dict, ...]:
    """Retrieve all genotypes datasets by species and population"""
    with conn.cursor(cursorclass=DictCursor) as cursor:
        cursor.execute(
            "SELECT gf.* FROM InbredSet AS iset INNER JOIN GenoFreeze AS gf "
            "ON iset.InbredSetId=gf.InbredSetId "
            "WHERE iset.SpeciesId=%(sid)s AND iset.InbredSetId=%(pid)s",
            {"sid": speciesid, "pid": populationid})
        return tuple(dict(row) for row in cursor.fetchall())

def geno_dataset_by_id(conn: mdb.Connection, dataset_id: int) -> dict:
    """Retrieve genotype dataset by ID"""
    with conn.cursor(cursorclass=DictCursor) as cursor:
        cursor.execute("SELECT * FROM GenoFreeze WHERE Id=%s",
                       (dataset_id,))
        return dict(cursor.fetchone())

def probeset_studies_by_species_and_population(
        conn: mdb.Connection,
        speciesid: int,
        populationid: int) -> tuple[dict, ...]:
    """Retrieve all probesets"""
    with conn.cursor(cursorclass=DictCursor) as cursor:
        cursor.execute(
            "SELECT pf.* FROM InbredSet AS iset INNER JOIN ProbeFreeze AS pf "
            "ON iset.InbredSetId=pf.InbredSetId "
            "WHERE iset.SpeciesId=%(sid)s AND iset.InbredSetId=%(pid)s",
            {"sid": speciesid, "pid": populationid})
        return tuple(dict(row) for row in cursor.fetchall())

def probeset_datasets_by_study(conn: mdb.Connection,
                               studyid: int) -> tuple[dict, ...]:
    """Retrieve all probeset databases by study."""
    with conn.cursor(cursorclass=DictCursor) as cursor:
        cursor.execute("SELECT * FROM ProbeSetFreeze WHERE ProbeFreezeId=%s",
                       (studyid,))
        return tuple(dict(row) for row in cursor.fetchall())

def probe_study_by_id(conn: mdb.Connection, studyid: int) -> dict:
    """Retrieve ProbeSet study by ID"""
    with conn.cursor(cursorclass=DictCursor) as cursor:
        cursor.execute("SELECT * FROM ProbeFreeze WHERE Id=%s",
                       (studyid,))
        return dict(cursor.fetchone())