from astroquery.mast import Observations
import numpy as np
from enum import Enum

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

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]