from astroquery.mast import Observations
from astropy import table
from astropy.io import fits
import numpy as np
from enum import Enum
from copy import deepcopy

class DataType(Enum):
    DataValidation = 1
    LightCurve = 2
    TargetPixel = 3

class ObservationSource(Enum):
    KEPLER = 1
    K2 = 2
    TESS = 3

# base from https://spacetelescope.github.io/mast_notebooks/notebooks/TESS/beginner_astroquery_dv/beginner_astroquery_dv.html
# modified for sector span
def parse_manifest(manifest):
    """
    Parse manifest and add back columns that are useful for TESS DV exploration.
    """
    results = deepcopy(manifest)
    filenames = []
    sources = []
    sector_range = []
    exts = []
    for i,f in enumerate(manifest['Local Path']):
        file_parts = np.array(np.unique(f.split(sep = '-')))
        try:
            with fits.open(f, mode="readonly") as hdu:
                if("SECTOR" in hdu[0].header):
                    sector_range.append(f"{hdu[0].header['SECTOR']}")
                elif("TTABLEID" in hdu[0].header):
                    sector_range.append(f"{hdu[0].header['TTABLEID']}")
                hdu.close()
        except:
            sector_range.append("-")

        if("TESS" in f):
            sources.append("TESS")
        elif("Kepler" in f):
            sources.append("Kepler")
        elif("K2" in f):
            sources.append("K2")

        path_parts = np.array(f.split(sep = '/'))
        filenames.append(path_parts[-1])
        exts.append(path_parts[-1][-8:])

    results.add_column(table.Column(name = "filename", data = filenames))
    results.add_column(table.Column(name = "source", data = sources))
    results.add_column(table.Column(name = "sector", data = sector_range))
    results.add_column(table.Column(name = "fileType", data = exts))
    results.add_column(table.Column(name = "index", data = np.arange(0,len(manifest))))

    return results

def getStarObservations(starName: str, sources: list[ObservationSource], sequences: list[list[int]] = []) -> str:
    obs = Observations.query_object(objectname=starName, radius="0 deg")

    obsWantedFilter = obs["dataproduct_type"] == "timeseries"

    if(len(sequences) > 1):
        obsWantedFilter &= ((obs['sequence_number'] >= max(sequences)) &
                  (obs['sequence_number'] <= min(sequences)))
    elif(len(sequences) == 1):
        obsWantedFilter &= (obs['sequence_number'] == sequences[0])

    obsSourceFilter = np.ma.MaskedArray(data=np.full(obsWantedFilter.shape, False),
                                        mask=False, fill_value=True)
    if(ObservationSource.KEPLER in sources):
        obsSourceFilter |= (obs['obs_collection'] == "Kepler")
    if(ObservationSource.K2 in sources):
        obsSourceFilter |= (obs['obs_collection'] == "K2")
    if(ObservationSource.TESS in sources):
        obsSourceFilter |= (obs['obs_collection'] == "TESS")

    obsWantedFilter &= obsSourceFilter

    return obs[obsWantedFilter]

def fitsFilenameFilterFunc(table, key_colnames):
    if(str(table["productFilename"]).lower().endswith(".fits") or
       str(table["productFilename"]).lower().endswith(".fit")):
        return True
    return False

def downloadStarProducts(obs_wanted, keplerCadences, k2Cadences):
    if(len(keplerCadences) != 2):
        raise Exception("keplerCadences needs to have length 2")
    if(len(k2Cadences) != 2):
        raise Exception("k2Cadences need to have length 2")

    dataProducts = Observations.get_product_list(obs_wanted)

    # LC:  TESS lightcurve
    # SLC: Kepler short cadence lightcurve
    # LLC: Kepler long cadence lightcurve
    productSubGroups = ["LC", "SLC"]
    tessProducts = Observations.filter_products(dataProducts, obs_collection="TESS",
                                                  productSubGroupDescription=["LC"])
    keplerSubGroups = []
    if(keplerCadences[0]):
        keplerSubGroups.append("SLC")
    if(keplerCadences[1]):
        keplerSubGroups.append("LLC")

    keplerProducts = Observations.filter_products(dataProducts, obs_collection="Kepler",
                                                  productSubGroupDescription=keplerSubGroups)

    k2SubGroups = []
    if(k2Cadences[0]):
        k2SubGroups.append("SLC")
    if(k2Cadences[1]):
        k2SubGroups.append("LLC")

    k2Products = Observations.filter_products(dataProducts, obs_collection="K2",
                                                  productSubGroupDescription=k2SubGroups)

    productsWanted = table.vstack([tessProducts, keplerProducts, k2Products])
    productsWanted = productsWanted.group_by("productFilename")
    productsWanted = productsWanted.groups.filter(fitsFilenameFilterFunc)

    filenames = Observations.download_products(productsWanted)
    return parse_manifest(filenames)