import polars as pl


def prepareQuery(queryData):

    local_type = queryData.get('queryType')  # KeyError protected -- returns None if the key is not defined

    #################
    if local_type=='occ_tables':
        return "SELECT name FROM sqlite_master WHERE type='table'"

    #################
    if local_type=='forma':
        try:
            data = queryData['data']
            dataNorm = queryData['dataNorm']
        except KeyError as err:
            raise KeyError('Missing required data for query type ' + local_type + ': ' + str(err))
        joinedQueryData = "'" + "' OR spec LIKE '".join(data) + "'"
        if len(dataNorm)==0:
            return f"SELECT spec AS forma, nocc AS occ, cod FROM form WHERE spec LIKE {joinedQueryData} ORDER BY idfor"
        else:
            joinedQueryDataNorm = "'" + "' OR norm LIKE '".join(dataNorm) + "'"
            return f"SELECT DISTINCT spec AS forma, nocc AS occ, cod FROM form WHERE (spec LIKE {joinedQueryData}) OR (norm LIKE {joinedQueryDataNorm}) ORDER BY idfor"

    ###################
    elif local_type=='lemma':
        try:
            data = queryData['data']
            dataNorm = queryData['dataNorm']
        except KeyError as err:
            raise KeyError('Missing required data for query type ' + local_type + ': ' + str(err))
        joinedQueryData = "'" + "' OR spec LIKE '".join(data) + "'"
        if len(dataNorm)==0:
            return f"SELECT spec AS lemma, cat AS cat_gr, omo AS disambiguatore, nocc AS occ, cod FROM lem WHERE spec LIKE {joinedQueryData} ORDER BY idlem"
        else:
            joinedQueryDataNorm = "'" + "' OR norm LIKE '".join(dataNorm) + "'"
            return f"SELECT DISTINCT spec AS lemma, cat AS cat_gr, omo AS disambiguatore, nocc AS occ, cod FROM lem WHERE (spec LIKE {joinedQueryData}) OR (norm LIKE {joinedQueryDataNorm}) ORDER BY idlem"

    ########################
    elif local_type=='lemmaForma':
        try:
            data = queryData['data']
            dataNorm = queryData['dataNorm']
        except KeyError as err:
            raise KeyError('Missing required data for query type ' + local_type + ': ' + str(err))
        joinedQueryData = "'" + "' OR form.spec LIKE '".join(data) + "'"
        if len(dataNorm)==0:
            return f"SELECT lem.spec AS lemma, lem.cat AS cat_gr, form.spec AS forma, lem.omo AS disambiguatore, pfl.nocc AS occ, lem.cod FROM pfl INNER JOIN form ON form.cod = pfl.forma INNER JOIN lem ON lem.cod != 0 AND lem.cod = pfl.lemma WHERE lem.spec LIKE {joinedQueryData} ORDER BY lem.idlem"
        else:
            joinedQueryData = "'" + "' OR lem.spec LIKE '".join(data) + "'"
            joinedQueryDataNorm = "'" + "' OR lem.norm LIKE '".join(dataNorm) + "'"
            return f"SELECT DISTINCT lem.spec AS lemma, lem.cat AS cat_gr, form.spec AS forma, lem.omo AS disambiguatore, pfl.nocc AS occ, lem.cod FROM pfl INNER JOIN form ON form.cod = pfl.forma INNER JOIN lem ON lem.cod = pfl.lemma WHERE (lem.spec LIKE {joinedQueryData}) OR (lem.norm LIKE {joinedQueryDataNorm}) ORDER BY lem.idlem"

    ########################
    elif local_type=='formaLemma':
        try:
            data = queryData['data']
            dataNorm = queryData['dataNorm']
        except KeyError as err:
            raise KeyError('Missing required data for query type ' + local_type + ': ' + str(err))
        joinedQueryData = "'" + "' OR form.spec LIKE '".join(data) + "'"
        if len(dataNorm)==0:
            return f"SELECT form.spec AS forma, lem.spec AS lemma, lem.cat AS cat_gr, lem.omo AS disambiguatore, pfl.nocc AS occ, form.cod FROM pfl INNER JOIN form ON form.cod = pfl.forma INNER JOIN lem ON lem.cod = pfl.lemma WHERE form.spec LIKE {joinedQueryData} ORDER BY form.idfor"
        else:
            joinedQueryDataNorm = "'" + "' OR form.norm LIKE '".join(dataNorm) + "'"
            return f"SELECT DISTINCT form.spec AS forma, lem.spec AS lemma, lem.cat AS cat_gr, lem.omo AS disambiguatore, pfl.nocc AS occ, form.cod FROM pfl INNER JOIN form ON form.cod = pfl.forma INNER JOIN lem ON lem.cod = pfl.lemma WHERE (form.spec LIKE {joinedQueryData}) OR (form.norm LIKE {joinedQueryDataNorm}) ORDER BY form.idfor"

    #################
    elif local_type=='pfl':
        try:
            codList = queryData['codList']
        except KeyError as err:
            raise KeyError('Missing required data for query type ' + local_type + ': ' + str(err))

        strlist = ",".join(str(c) for c in codList)
        return f"SELECT DISTINCT lemma as codLemma, forma as codForma FROM pfl WHERE lemma IN ({strlist})"

    ###################
    elif local_type=='texts':
        return complexQueryTexts
    
    ###################
    elif local_type=='co-occurrences':
        return complexQueryCooccurrences

    ######################
    elif local_type=='bib':
        try:
            row = queryData['row']
            sigla = row['sigla']
        except KeyError as err:
            raise KeyError('Missing required data for query type ' + local_type + ': ' + str(err))
        return f"SELECT [Anno iniziale], [Anno finale], [Data codificata], [Titolo Abbreviato], [Autore], [Titolo], [Curatore], [Data descrittiva], [Area generica], [Area specifica], [Genere], [Forma], [Tipo], IQ FROM datibib WHERE Sigla='{sigla}'"
    
    #################
    elif local_type=='bibAlt':
        try:
            siglaSet = queryData['siglaSet']
        except KeyError as err:
            raise KeyError('Missing required data for query type ' + local_type + ': ' + str(err))
        siglaStr = "'" + "','".join(siglaSet) + "'"
        return f"SELECT Sigla, [Anno iniziale], [Anno finale], [Data codificata], [Titolo Abbreviato], [Autore], [Titolo], [Curatore], [Data descrittiva], [Area generica], [Area specifica], [Genere], [Forma], [Tipo], IQ FROM datibib WHERE Sigla IN ({siglaStr})"

    #################
    elif local_type=='rif':
        try:
            row = queryData['row']
            numorg = row['numorg']
            ntx = row['ntx']
        except:
            return None
        return f"SELECT head AS Rif_organico, full AS Rif_completo FROM org WHERE (indice='{numorg}' AND ntx='{ntx}')"
    
    #################
    elif local_type=='rifAlt':
        try:
            coordsSet = queryData['coordsSet']
        except KeyError as err:
            raise KeyError('Missing required data for query type ' + local_type + ': ' + str(err))
        
        subQueries = []
        for coords in coordsSet:
            try:
                numorg = coords[0]
                ntx = coords[1]
            except IndexError as err:
                raise KeyError('Incomplete required data for query type ' + local_type + ': ' + str(err))
            subQueries.append( f"SELECT indice AS numorg, ntx, head AS Rif_organico, full AS Rif_completo FROM org WHERE (indice='{numorg}' AND ntx='{ntx}')" )
                
        return ' UNION ALL '.join(subQueries)

    #################
    elif local_type=='highlight':
        try:
            row = queryData['row']
            col = queryData['col']
        except:
            return None
        return f"SELECT spec as highlight FROM form WHERE cod={row[col]}"
    
    #################
    elif local_type =='singlecontext':
        try:
            subtype = queryData['querySubtype']
            table = queryData['table']
            parole = queryData['parole']
            periodi = queryData['periodi']
            brani = queryData['brani']
            ntxlocal = queryData['ntxlocal']
            mappalocal = queryData['mappalocal']
            periodlocal = queryData['periodlocal']
            numbranolocal = queryData['numbranolocal']
        except:
            return None
        if subtype == 'parole':
            return f"SELECT tab.pitxt, tab.elemlen FROM {table} AS tab WHERE tab.ntx = {ntxlocal} AND tab.mappa <= {mappalocal+int(parole/2)} AND tab.mappa >= {mappalocal-int(parole/2)}"
        elif subtype == 'periodi':
            return f"SELECT piniz, pfin FROM periodi WHERE ntx = {ntxlocal} AND numperiod <= {periodlocal+int(periodi/2)} AND numperiod >= {periodlocal-int(periodi/2)}"
        elif subtype == 'brani':
            return f"SELECT piniz, pfin FROM linkbase WHERE {ntxlocal} = ntx AND tipo = 2 AND id BETWEEN {numbranolocal-int(brani/2)} AND {numbranolocal+int(brani/2)}"
    #################
    elif local_type =='links':
        try:
            subtype = queryData['querySubtype']
            ntxlocal = queryData['ntxlocal']
            pinizlocal = queryData['pinizlocal']
            pfinlocal = queryData['pfinlocal']
            pitxtlocal = queryData['pitxtlocal']
        except:
            return None   
        if subtype == 'nota':
            return f"SELECT ta.ntx, ta.id, ta.piniz, ta.pfin, tb.mappain, tb.mappafin FROM linkbase AS tb INNER JOIN linknoteass AS ta ON tb.ntx = ta.ntx AND tb.id = ta.id WHERE (((tb.tipo= 1) AND (tb.ntx = {ntxlocal})) AND ((tb.piniz BETWEEN {pinizlocal} AND {pfinlocal}) OR ({pitxtlocal} BETWEEN tb.piniz AND tb.pfin)))"
        if subtype == 'testo_associato':
            return f"SELECT ta.ntx, ta.id, ta.piniz, ta.pfin, tb.mappain, tb.mappafin FROM linkbase AS tb INNER JOIN linknoteass AS ta ON tb.ntx = ta.ntx AND tb.id = ta.id WHERE (((tb.tipo= 2) AND (tb.ntx = {ntxlocal})) AND ((tb.piniz BETWEEN {pinizlocal} AND {pfinlocal}) OR ({pitxtlocal} BETWEEN tb.piniz AND tb.pfin)))"
    #####
    else:
        raise ValueError('Unrecognized query type: ' + local_type)


def complexQueryTexts(connection, queryData):
    try:
        codList = queryData['codList']
        table = queryData['table']
        subtype = queryData['querySubtype']
        formCodList = queryData.get('formCodList')  # KeyError-safe (None if absent)
    except KeyError as err:
        raise KeyError(f'Missing required data for query type {queryData.get('queryType')} : {str(err)}')

    strCodList = ",".join(str(c) for c in codList)

    # Main query, verified to be fast!
    mainQueryString = f"""
    SELECT tab.cod, tab.ntx, tab.pitxt, tab.elemlen, tab.mappa, tab.numperiod, tab.links, tab.numorg, 
           intbib.sigla, tab.vol, tab.pag, tab.riga, tab.col, tab.tipostanza, tab.stanza, tab.verso, 
           tab.numbrano, lem.spec AS lemma, lem.cat AS cat_gr, lem.omo AS disambiguatore 
    FROM {table} AS tab 
    INNER JOIN intbib ON tab.ntx = intbib.ntx 
    INNER JOIN lem ON tab.indlem = lem.cod
    """

    if subtype == 0:
        condition = f"WHERE tab.cod IN ({strCodList})"
    elif subtype == 1:
        condition = f"WHERE tab.indlem IN ({strCodList})"
    elif subtype == 2:
        if formCodList is None:
            return None
        strFormCodList = ",".join(str(c) for c in formCodList)
        condition = f" WHERE tab.indlem IN ({strCodList}) OR (tab.indlem = 0 AND tab.cod IN ({strFormCodList}))"

    mainQueryString = f'{mainQueryString} {condition}'

    # This value can be changed to change multiple contexts' width. Default value for Gatto is parole=31 #
    parole = 31
    # C'è la possibilità di scegliere periodi invece che parole, ma per il momento è disabilitata
    createTempTable = f'CREATE TEMPORARY TABLE stuff AS {mainQueryString}'
    mainQuery = 'SELECT * from stuff'
    addQuery1 = f'SELECT stuff.ntx, stuff.mappa, tab.pitxt AS piniz FROM stuff LEFT JOIN Occ00001 AS tab ON tab.ntx=stuff.ntx AND tab.mappa=stuff.mappa-{int(parole / 2)}'
    addQuery2 = f'SELECT stuff.ntx, stuff.mappa, tab.pitxt AS pfin FROM stuff LEFT JOIN Occ00001 AS tab ON tab.ntx=stuff.ntx AND tab.mappa=stuff.mappa+{int(parole/2)}'
    addQuery3 = 'SELECT stuff.ntx, stuff.numperiod, periodi.piniz AS backup_piniz, periodi.pfin AS backup_pfin FROM stuff, periodi WHERE stuff.ntx = periodi.ntx AND stuff.numperiod = periodi.numperiod'

    # Start communication with DB
    cursor = connection.cursor()
    cursor.execute(createTempTable)

    results = pl.read_sql(mainQuery, connection)
    results_add1 = pl.read_sql(addQuery1, connection)
    results_add2 = pl.read_sql(addQuery2, connection)
    results_add3 = pl.read_sql(addQuery3, connection)

    results = results.with_columns([
        results_add1['piniz'],
        results_add2['pfin'],
        results_add3['backup_piniz'],
        results_add3['backup_pfin']
    ])

    return results


def complexQueryCooccurrences(connection, queryData):
    try:
        occurrences = queryData['occurrences']
        table = queryData['table']
        intervallo = queryData['intervallo']
        periodo = queryData.get('periodo', 0) # Unused for the moment
        ordinate = queryData.get('ordinate', 0) # Unused for the moment
    except KeyError as err:
        raise KeyError('Missing required data for query: ' + str(err))

    preMainQueryString = f"""
    SELECT tab.cod, tab.ntx, tab.pitxt, tab.elemlen, tab.mappa, tab.numperiod, tab.links, tab.numorg, 
           intbib.sigla, tab.vol, tab.pag, tab.riga, tab.col, tab.tipostanza, tab.stanza, tab.verso, 
           tab.numbrano, lem.spec AS lemma, lem.cat AS cat_gr, lem.omo AS disambiguatore 
    FROM {table} AS tab 
    INNER JOIN intbib ON tab.ntx = intbib.ntx 
    INNER JOIN lem ON tab.indlem = lem.cod
    """

    pitxtList = ['pitxt']
    elemlenList = ['elemlen']

    for index, occ in enumerate(occurrences):
        try:
            subtype = occ['querySubtype']
            codList = occ['codList']
            formCodList = occ.get('formCodList')
        except KeyError as err:
            raise KeyError('Missing required data for query: ' + str(err))

        strCodList = ",".join(str(c) for c in codList)

        if subtype == 0:
            condition = f" WHERE tab.cod IN ({strCodList})"
        elif subtype == 1:
            condition = f" WHERE tab.indlem IN ({strCodList})"
        elif subtype == 2:
            if formCodList is None:
                return None
            strFormCodList = ",".join(str(c) for c in formCodList)
            condition = f" WHERE tab.indlem IN ({strCodList}) OR (tab.indlem = 0 AND tab.cod IN ({strFormCodList}))"

        mainQueryString = f'{preMainQueryString} {condition}'

        if index == 0:
            resTable = 'tempOcc_' + str(index)
            cursor = connection.cursor()
            cursor.execute(f'CREATE TEMPORARY TABLE {resTable} AS {mainQueryString}')
            cursor.execute(f'CREATE INDEX aa_{index} ON {resTable} (ntx, mappa)')
        else:
            cursor.execute(f'CREATE TEMPORARY TABLE tempOccB AS {mainQueryString}')
            cursor.execute(f'CREATE INDEX bb ON tempOccB (ntx, mappa)')

            oldTable = resTable
            resTable = 'tempOcc_' + str(index)
            cursor.execute(f"""
            CREATE TEMPORARY TABLE {resTable} AS 
            SELECT tabA.cod, tabA.ntx, tabA.{" tabA.".join(pitxtList)}, tabA.{" tabA.".join(elemlenList)}, 
                   tabA.mappa, tabA.numperiod, tabA.links, tabA.numorg, tabA.sigla, tabA.vol, tabA.pag, 
                   tabA.riga, tabA.col, tabA.tipostanza, tabA.stanza, tabA.verso, tabA.numbrano, 
                   tabA.lemma, tabA.cat_gr, tabA.disambiguatore, tabB.ntx AS ntx2, tabB.mappa AS mappa2, 
                   tabB.pitxt as pitxt_{index}, tabB.elemlen as elemlen_{index} 
            FROM {oldTable} AS tabA, tempOccB AS tabB 
            WHERE tabA.ntx=tabB.ntx AND tabA.mappa BETWEEN tabB.mappa-{intervallo} AND tabB.mappa+{intervallo} AND tabA.mappa != tabB.mappa
            """)
            cursor.execute(f'CREATE INDEX aa_{index} ON {resTable} (ntx, mappa)')
            cursor.execute(f'DROP TABLE {oldTable}')
            pitxtList.append(f'pitxt_{index}')
            elemlenList.append(f'elemlen_{index}')

    results = pl.read_sql(f'SELECT * FROM {resTable}', connection)

    parole = 31
    queryPiniz = f'SELECT stuff.ntx, stuff.mappa, tab.pitxt AS piniz FROM {resTable} AS stuff LEFT JOIN {table} AS tab ON tab.ntx=stuff.ntx AND tab.mappa=stuff.mappa-{int(parole/2)}'
    queryPfin = f'SELECT stuff.ntx, stuff.mappa, tab.pitxt AS pfin FROM {resTable} AS stuff LEFT JOIN {table} AS tab ON tab.ntx=stuff.ntx AND tab.mappa=stuff.mappa+{int(parole/2)}'
    queryPeriodi = f'SELECT stuff.ntx, stuff.numperiod, periodi.piniz AS backup_piniz, periodi.pfin AS backup_pfin FROM {resTable} AS stuff, periodi WHERE stuff.ntx = periodi.ntx AND stuff.numperiod = periodi.numperiod'

    resultsPiniz = pl.read_sql(queryPiniz, connection)
    resultsPfin = pl.read_sql(queryPfin, connection)
    resultsPeriodi = pl.read_sql(queryPeriodi, connection)

    results = results.with_columns([
        resultsPiniz['piniz'],
        resultsPfin['pfin'],
        resultsPeriodi['backup_piniz'],
        resultsPeriodi['backup_pfin']
    ])

    return results