"""Functions for handling genotypes."""
from typing import Optional

import MySQLdb as mdb
from MySQLdb.cursors import DictCursor

from uploader.db_utils import debug_query

def genocode_by_population(
        conn: mdb.Connection, population_id: int) -> tuple[dict, ...]:
    """Get the allele/genotype codes."""
    with conn.cursor(cursorclass=DictCursor) as cursor:
        cursor.execute("SELECT * FROM GenoCode WHERE InbredSetId=%s",
                       (population_id,))
        return tuple(dict(item) for item in cursor.fetchall())


def genotype_markers_count(conn: mdb.Connection, species_id: int) -> int:
    """Find the total count of the genotype markers for a species."""
    with conn.cursor(cursorclass=DictCursor) as cursor:
        cursor.execute(
            "SELECT COUNT(Name) AS markers_count FROM Geno WHERE SpeciesId=%s",
            (species_id,))
        return int(cursor.fetchone()["markers_count"])


def genotype_markers(
        conn: mdb.Connection,
        species_id: int,
        offset: int = 0,
        limit: Optional[int] = None
) -> tuple[dict, ...]:
    """Retrieve markers from the database."""
    _query = "SELECT * FROM Geno WHERE SpeciesId=%s"
    if bool(limit) and limit > 0:
        _query = _query + f" LIMIT {limit} OFFSET {offset}"

    with conn.cursor(cursorclass=DictCursor) as cursor:
        cursor.execute(_query, (species_id,))
        debug_query(cursor)
        return tuple(dict(row) for row in cursor.fetchall())


def genotype_datasets(
        conn: mdb.Connection,
        species_id: int,
        population_id: int
) -> tuple[dict, ...]:
    """Retrieve genotype datasets from the database."""
    with conn.cursor(cursorclass=DictCursor) as cursor:
        cursor.execute(
            "SELECT gf.* FROM Species AS s INNER JOIN InbredSet AS iset "
            "ON s.Id=iset.SpeciesId INNER JOIN GenoFreeze AS gf "
            "ON iset.Id=gf.InbredSetId "
            "WHERE s.Id=%s AND iset.Id=%s",
            (species_id, population_id))
        return tuple(dict(row) for row in cursor.fetchall())