from astropy.coordinates import SkyCoord
from astropy.io.votable import parse
import requests
import io
import sys

class SingleSource():
    """
    Small class to use with the ARI single source service. In order to
    make it somewhat VO-compliant and interactive with PyVO and Astropy,
    result are returned on VOTable format by default
    """
    base_url = "https://gaia.ari.uni-heidelberg.de/singlesource/"
    catalogues = ["gaiadr3", "gaiaedr3", "gaiadr2", "gaiadr1", "tgas"]
    session = requests.session()
    headers = {'Connection': 'close'}
    
    def __get(self, url, params):
        """Runs an HTTP GET query"""
        return self.session.get(url, params=params, headers=self.headers)

    def __post(self, url, files, params):
        """Runs an HTTP POST query"""
        return self.session.post(url, files=files, params=params, headers=self.headers)

    def search_by_name(self, catalogue, obj, more=None):
        """
        Searches by a given celestial object

        Parameters:
            - catalogue: valid catalogue: DR3, EDR3, DR2, DR1, TGAS
            - obj: name ob the object
            - more: whether you'd like more than one result
        """
        if catalogue not in self.catalogues:
            print("Incorrect catalogue, valid values are: " + str(self.catalogues))
            sys.exit()
        params = {'obj': obj, 'f': 'votable'}
        if more is not None:
            params['v'] = True
        response = self.__get(self.base_url + catalogue, params)
        if response.status_code == requests.codes.ok:
            votable = io.BytesIO(response.text.encode())
            return parse(votable) 
        else:
            return None

    def search_by_position(self, catalogue, coords, more=None):
        """
        Searches by a given position
        
        Parameters:
            - catalogue: valid catalogue: DR3, EDR3, DR2, DR1, TGAS
            - coords, either an ICRS SkyCoord or a (ra,dec) tuple
            - more: whether you'd like more than one result
        """
        if catalogue not in catalogues:
            print("Incorrect catalogue, valid values are: " + str(self.catalogues))
            sys.exit()
        params = {'ra' : 0, 'dec' : 0, 'f': 'votable'}
        if type(coords) is SkyCoord:
            params['ra'] = coords.ra.degree
            params['dec'] = coords.dec.degree
        elif type(coords) is Tuple and len(coords) == 2:
            params['ra'] = coords[0]
            params['dec'] = coords[1]
        if more is not None:
            params['v'] = True
        response = self.__get(self.base_url + catalogue, params)
        if response.status_code == requests.codes.ok:
            votable = io.BytesIO(response.text.encode())
            return parse(votable)
        else:
            return None

    def search_by_id(self, catalogue, sourceid, more=None):
        """
        Searches an item by identifier
        
        Parameters:
            - catalogue: valid catalogue: DR3, EDR3, DR2, DR1, TGAS
            - id: numeric identifier
            - more: whether you'd like more than one result
        """
        if catalogue not in self.catalogues:
            print("Incorrect catalogue, valid values are: " + str(self.catalogues))
            sys.exit()
        params={'id' : sourceid, 'f': 'votable'}
        if more is not None:
            params['v'] = True
        response = self.__get(self.base_url + catalogue, params)
        if response.status_code == requests.codes.ok:
            votable = io.BytesIO(response.text.encode())
            return parse(votable)
        else:
            return None

    def search_by_file(self, catalogue, filename):
        """
        Searches with the information retrieved from a CSV file, which
        should consisting of source ID, positions or
        coordinates per each line.
        
        Parameters:
            - catalogue: valid catalogue: DR3, EDR3, DR2, DR1, TGAS
            - filename: name of the file
        """ 
        if catalogue not in self.catalogues:
            print("Incorrect catalogue, valid values are: " + str(self.catalogues))
            sys.exit()
        if filename.split(".")[-1] != "csv":
            print("File must be in CSV format")
            sys.exit()
        files = {'positions': open(filename, 'rb')}
        params = {'f': 'votable'}
        response = self.__post(self.base_url + catalogue, files, params)
        if response.status_code == requests.codes.ok:
            votable = io.BytesIO(response.text.encode())
            return parse(votable) 
        else:
            return None