# Licensed under a 3-clause BSD style license - see LICENSE.rst

import os
import pytest
import tempfile
from unittest.mock import patch, PropertyMock
from astropy.coordinates import SkyCoord
from astropy.table import Table
from astropy.io import votable
from pyvo.dal import TAPResults
import astropy.units as u

from astroquery.heasarc import Heasarc, HeasarcClass
from astroquery.exceptions import InvalidQueryError

try:
    # Both boto3, botocore and moto are optional dependencies,
    # but the former 2 are dependencies of the latter, so it's enough
    # to handle them with one variable
    import boto3
    from moto import mock_aws

    DO_AWS_S3 = True
except ImportError:
    DO_AWS_S3 = False


OBJ_LIST = [
    "182d38m08.64s 39d24m21.06s",
    SkyCoord(l=155.0771, b=75.0631, unit=(u.deg, u.deg), frame="galactic"),
]

SIZE_LIST = [2 * u.arcmin, "0d2m0s"]


# used in the tap fixture
class MockTap:
    class vColumn:
        def __init__(self, name, desc, unit=None):
            self.name = name
            self.description = desc
            self.unit = unit

    class vTable:
        def __init__(self, desc, cols=[]):
            self.description = desc
            self.columns = cols

    cols = []
    for icol in [1, 2, 3]:
        cols.append(vColumn(f'col-{icol}', f'desc-{icol}'))

    tables = {
        'name-1': vTable('description-1 xmm', cols),
        'name-2': vTable('description-2 chandra', cols),
        'TAPname': None
    }

    def search(self, query, language='ADQL', maxrec=1000, uploads=None):
        if "SELECT COUNT(*) FROM" in query:
            return TAPResults(votable.from_table(Table([[3055]], names=['count'],
                                                       dtype=['int64'])))
        else:
            return MockResult()


class MockResult:
    def to_table(self):
        return Table({'value': ['1.5', '1.2', '-0.3']})


@pytest.fixture
def mock_tap():
    with patch('astroquery.heasarc.core.HeasarcClass.tap', new_callable=PropertyMock) as tap:
        tap.return_value = MockTap()
        yield tap


@pytest.fixture
def mock_default_cols():
    with patch('astroquery.heasarc.core.HeasarcClass._get_default_columns') as get_cols:
        get_cols.return_value = ['col-3', 'col-2']
        yield get_cols


@pytest.fixture
def mock_meta():
    with patch('astroquery.heasarc.core.HeasarcClass._meta', new_callable=PropertyMock) as meta:
        meta.return_value = Table(dict(
            table=['tab1', 'tab2', 'tab1', 'tab1'],
            par=['p1', 'p2', 'p3', ''],
            value=[1.2, 1.6, 2.0, 3.0]
        ))
        yield meta


@pytest.mark.parametrize("coordinates", OBJ_LIST)
@pytest.mark.parametrize("radius", SIZE_LIST)
@pytest.mark.parametrize("offset", [True, False])
def test_query_region_cone(coordinates, radius, offset):
    # use columns='*' to avoid remote call to obtain the default columns
    query = Heasarc.query_region(
        coordinates,
        catalog="suzamaster",
        spatial="cone",
        radius=radius,
        columns="*",
        get_query_payload=True,
        add_offset=offset,
    )

    # We don't fully float compare in this string, there are slight
    # differences due to the name-coordinate resolution and conversions
    if offset:
        distance_text = ",DISTANCE(POINT('ICRS',ra,dec), POINT('ICRS',182.63"
    else:
        distance_text = ""
    assert (f"SELECT *{distance_text}") in query
    assert (
        "FROM suzamaster WHERE CONTAINS(POINT('ICRS',ra,dec),"
        "CIRCLE('ICRS',182.63" in query
    )
    assert ",39.40" in query
    assert ",0.0333" in query


@pytest.mark.parametrize("coordinates", OBJ_LIST)
@pytest.mark.parametrize("width", SIZE_LIST)
def test_query_region_box(coordinates, width):
    query = Heasarc.query_region(
        coordinates,
        catalog="suzamaster",
        spatial="box",
        width=2 * u.arcmin,
        columns="*",
        get_query_payload=True,
    )

    assert (
        "SELECT * FROM suzamaster WHERE CONTAINS(POINT('ICRS',ra,dec),"
        "BOX('ICRS',182.63" in query
    )
    assert ",39.40" in query
    assert ",0.0333" in query


poly1 = [
    SkyCoord(ra=10.1 * u.deg, dec=10.1 * u.deg),
    SkyCoord(ra=10.0 * u.deg, dec=10.1 * u.deg),
    SkyCoord(ra=10.0 * u.deg, dec=10.0 * u.deg),
]
poly2 = [
    (10.1 * u.deg, 10.1 * u.deg),
    (10.0 * u.deg, 10.1 * u.deg),
    (10.0 * u.deg, 10.0 * u.deg),
]


@pytest.mark.parametrize("polygon", [poly1, poly2])
def test_query_region_polygon(polygon):
    # position is not used for polygon
    query1 = Heasarc.query_region(
        catalog="suzamaster",
        spatial="polygon",
        polygon=polygon,
        columns="*",
        get_query_payload=True,
    )
    query2 = Heasarc.query_region(
        "ngc4151",
        catalog="suzamaster",
        spatial="polygon",
        polygon=polygon,
        columns="*",
        get_query_payload=True,
    )

    assert query1 == query2
    assert query1 == (
        "SELECT * FROM suzamaster "
        "WHERE CONTAINS(POINT('ICRS',ra,dec),POLYGON('ICRS',"
        "10.1,10.1,10.0,10.1,10.0,10.0))=1"
    )


def test_query_region_polygon_no_unit():
    # position is not used for polygon
    poly = [
        (10.1, 10.1),
        (10.0, 10.1),
        (10.0, 10.0),
    ]
    with pytest.warns(UserWarning, match="Polygon endpoints are being interpreted as"):
        query = Heasarc.query_region(
            catalog="suzamaster",
            spatial="polygon",
            polygon=poly,
            columns="*",
            get_query_payload=True,
        )

    assert query == (
        "SELECT * FROM suzamaster "
        "WHERE CONTAINS(POINT('ICRS',ra,dec),POLYGON('ICRS',"
        "10.1,10.1,10.0,10.1,10.0,10.0))=1"
    )


def test_query_allsky():
    query1 = Heasarc.query_region(
        catalog="suzamaster", spatial="all-sky", columns="*",
        get_query_payload=True
    )
    query2 = Heasarc.query_region(
        "m31",
        catalog="suzamaster",
        spatial="all-sky",
        columns="*",
        get_query_payload=True,
    )

    assert query1 == query2 == "SELECT * FROM suzamaster"


@pytest.mark.parametrize("spatial", ["space", "invalid"])
def test_spatial_invalid(spatial):
    with pytest.raises(ValueError):
        Heasarc.query_region(
            OBJ_LIST[0], catalog="invalid_spatial", columns="*", spatial=spatial
        )


def test_spatial_cone_no_position():
    with pytest.raises(InvalidQueryError):
        Heasarc.query_region(catalog="xmmmaster", columns="*", spatial="cone")


def test_no_catalog():
    with pytest.raises(InvalidQueryError):
        # OBJ_LIST[0] and radius added to avoid a remote call
        Heasarc.query_region(
            OBJ_LIST[0], spatial="cone", columns="*", radius="2arcmin")


def test__query_execute_no_catalog():
    with pytest.raises(InvalidQueryError):
        # OBJ_LIST[0] and radius added to avoid a remote call
        Heasarc._query_execute(None)


def test_parse_constraints_no_filter():
    assert Heasarc._parse_constraints(column_filters=None) == []
    assert Heasarc._parse_constraints(column_filters={}) == []


def test_parse_constraints_range():
    constraints = Heasarc._parse_constraints(column_filters={"flux": (1e-12, 1e-10)})
    assert constraints == ["flux BETWEEN 1e-12 AND 1e-10"]


def test_parse_constraints_eq_float():
    constraints = Heasarc._parse_constraints(column_filters={"flux": 1.2})
    assert constraints == ["flux = 1.2"]


def test_parse_constraints_eq_str():
    constraints = Heasarc._parse_constraints(column_filters={"flux": "1.2"})
    assert constraints == ["flux = '1.2'"]


def test_parse_constraints_cmp_float():
    constraints = Heasarc._parse_constraints(column_filters={"flux": ('>', 1.2)})
    assert constraints == ["flux > 1.2"]


def test_parse_constraints_cmp_float_2():
    constraints = Heasarc._parse_constraints(column_filters={"flux": ('>', 1.2), "magnitude": ('<=', 15)})
    assert constraints == ["flux > 1.2", "magnitude <= 15"]


def test_parse_constraints_list():
    constraints = Heasarc._parse_constraints(column_filters={"flux": [1.2, 2.3, 3.4]})
    assert constraints == ["flux IN (1.2, 2.3, 3.4)"]


def test_query_region_no_filter():
    query = Heasarc.query_region(
        catalog="suzamaster",
        spatial="all-sky",
        column_filters={},
        columns="*",
        get_query_payload=True,
    )
    assert query == "SELECT * FROM suzamaster"


def test_query_region_filter_range():
    query = Heasarc.query_region(
        catalog="suzamaster",
        spatial="all-sky",
        column_filters={"flux": (1e-12, 1e-10)},
        columns="*",
        get_query_payload=True,
    )
    assert query == "SELECT * FROM suzamaster WHERE flux BETWEEN 1e-12 AND 1e-10"


def test_query_region_filter_eq_float():
    query = Heasarc.query_region(
        catalog="suzamaster",
        spatial="all-sky",
        column_filters={"flux": 1.2},
        columns="*",
        get_query_payload=True,
    )
    assert query == "SELECT * FROM suzamaster WHERE flux = 1.2"


def test_query_region_filter_eq_str():
    query = Heasarc.query_region(
        catalog="suzamaster",
        spatial="all-sky",
        column_filters={"flux": "1.2"},
        columns="*",
        get_query_payload=True,
    )
    assert query == "SELECT * FROM suzamaster WHERE flux = '1.2'"


def test_query_region_filter_cmp_float():
    query = Heasarc.query_region(
        catalog="suzamaster",
        spatial="all-sky",
        column_filters={"flux": ('>', 1.2)},
        columns="*",
        get_query_payload=True,
    )
    assert query == "SELECT * FROM suzamaster WHERE flux > 1.2"


def test_query_region_filter_cmp_float_2():
    query = Heasarc.query_region(
        catalog="suzamaster",
        spatial="all-sky",
        column_filters={"flux": ('>', 1.2), "magnitude": ('<=', 15)},
        columns="*",
        get_query_payload=True,
    )
    assert query == ("SELECT * FROM suzamaster WHERE flux > 1.2 "
                     "AND magnitude <= 15")


def test_query_region_filter_list():
    query = Heasarc.query_region(
        catalog="suzamaster",
        spatial="all-sky",
        column_filters={"flux": [1.2, 2.3, 3.4]},
        columns="*",
        get_query_payload=True,
    )
    assert query == "SELECT * FROM suzamaster WHERE flux IN (1.2, 2.3, 3.4)"


@pytest.mark.parametrize("coordinates", OBJ_LIST)
def test_query_region_cone_with_filter(coordinates):
    # use columns='*' to avoid remote call to obtain the default columns
    query = Heasarc.query_region(
        coordinates,
        catalog="suzamaster",
        spatial="cone",
        radius=2 * u.arcmin,
        columns="*",
        get_query_payload=True,
        column_filters={"flux": (1e-12, 1e-10)},
    )
    assert ("SELECT *") in query
    assert (
        "FROM suzamaster WHERE CONTAINS(POINT('ICRS',ra,dec),"
        "CIRCLE('ICRS',182.63" in query
    )
    assert ",39.40" in query
    assert ",0.0333" in query
    assert "AND flux BETWEEN 1e-12 AND 1e-10" in query


def test__query_execute_none_where():
    query = Heasarc._query_execute(
        catalog="suzamaster",
        columns="*",
        get_query_payload=True,
    )
    assert query == ("SELECT * FROM suzamaster")


def test__query_execute_extra_where():
    query = Heasarc._query_execute(
        catalog="suzamaster",
        where=" EXTRA",
        columns="*",
        get_query_payload=True,
    )
    assert query == ("SELECT * FROM suzamaster EXTRA")


def test__query_execute_add_row():
    query1 = Heasarc._query_execute(
        catalog="suzamaster",
        where="",
        columns="col1, col2",
        get_query_payload=True,
    )
    query2 = Heasarc._query_execute(
        catalog="suzamaster",
        where=None,
        columns="col1, col2",
        get_query_payload=True,
    )
    assert query1 == query2 == ("SELECT col1, col2, __row FROM suzamaster")


def test__query_execute_extra_space():
    query1 = Heasarc._query_execute(
        catalog="suzamaster",
        where="WHERE EXTRA",
        columns="*",
        get_query_payload=True,
    )

    query2 = Heasarc._query_execute(
        catalog="suzamaster",
        where=" WHERE EXTRA",
        columns="*",
        get_query_payload=True,
    )
    assert query1 == query2 == ("SELECT * FROM suzamaster WHERE EXTRA")


def test_query_execute_columns1(mock_tap, mock_default_cols):
    query = Heasarc._query_execute(
        catalog="suzamaster",
        where="WHERE EXTRA",
        columns=None,
        get_query_payload=True,
    )
    assert query == ("SELECT col-3, col-2, __row FROM suzamaster WHERE EXTRA")


def test_query_execute_columns2(mock_tap, mock_default_cols):
    _ = Heasarc._query_execute(
        catalog="suzamaster",
        where="WHERE EXTRA",
        columns='*'
    )
    assert Heasarc._last_catalog_name == "suzamaster"
    # reset last result to avoid interference with other tests
    Heasarc._last_result = None


def test_tap_def():
    # Use a new HeasarcClass object
    Heasarc = HeasarcClass()
    assert Heasarc._tap is None
    Heasarc.tap
    assert Heasarc._tap is not None


def test_meta_def():
    # Use a new HeasarcClass object
    Heasarc = HeasarcClass()
    assert Heasarc._meta_info is None
    with patch('astroquery.heasarc.core.HeasarcClass.query_tap') as mock_query_tap:
        mock_query_tap.return_value = MockResult()
        meta = Heasarc._meta
        assert meta['value'][0] == 1.5


def test__get_default_columns(mock_meta):
    # Use a new HeasarcClass object
    Heasarc = HeasarcClass()
    cols = Heasarc._get_default_columns('tab1')
    assert list(cols) == ['p1', 'p3']


def test__get_default_radius(mock_meta):
    # Use a new HeasarcClass object
    Heasarc = HeasarcClass()
    radius = Heasarc.get_default_radius('tab1')
    assert radius.value == 3.0


def test__set_session():
    with pytest.raises(ValueError):
        Heasarc._set_session('not session object')


def test__list_catalogs(mock_tap):
    catalogs = Heasarc.list_catalogs()
    assert list(catalogs['name']) == [
        lab for lab in MockTap().tables.keys() if 'TAP' not in lab
    ]
    assert list(catalogs['description']) == [
        desc.description for lab, desc in MockTap().tables.items() if 'TAP' not in lab
    ]


def test_list_catalogs_keywords_non_str():
    with pytest.raises(ValueError, match="non-str found in keywords elements"):
        Heasarc.list_catalogs(keywords=12)


def test_list_catalogs_keywords_list_non_str():
    with pytest.raises(ValueError, match="non-str found in keywords elements"):
        Heasarc.list_catalogs(keywords=['x-ray', 12])


def test__list_catalogs_keywords(mock_tap):
    catalogs = Heasarc.list_catalogs(keywords=['xmm'])
    assert list(catalogs['name']) == [
        lab for lab, desc in MockTap().tables.items() if 'TAP' not in lab and 'xmm' in desc.description.lower()
    ]


def test__list_columns__missing_table(mock_tap):
    with pytest.raises(ValueError, match="not available as a public catalog"):
        Heasarc.list_columns(catalog_name='missing-table')


def test__list_columns(mock_tap, mock_default_cols):
    cols = Heasarc.list_columns(catalog_name='name-1')
    assert list(cols['name']) == ['col-2', 'col-3']
    assert list(cols['description']) == ['desc-2', 'desc-3']


def test_locate_data():
    with pytest.raises(ValueError, match="query_result is None"):
        Heasarc.locate_data()

    with pytest.raises(
        TypeError, match=(
            "query_result need to be an astropy.table.Table or astropy.table.Row"
        )
    ):
        Heasarc.locate_data([1, 2])

    with pytest.raises(ValueError, match="No __row column found"):
        Heasarc.locate_data(Table({"id": [1, 2, 3.0]}), catalog_name="xray")


def test_locate_data_row():
    table = Table({"id": [1, 2, 3.0]})

    # we cannot do full call as this is not remote,
    # but if we check that we error on __row not input type
    with pytest.raises(ValueError, match="No __row column found"):
        Heasarc.locate_data(table[0], catalog_name="xray")

    with pytest.raises(ValueError, match="No __row column found"):
        Heasarc.locate_data(table[0:2], catalog_name="xray")


def test__guess_host_default():
    # Use a new HeasarcClass object
    assert Heasarc._guess_host(host=None) == 'heasarc'


@pytest.mark.parametrize("host", ["heasarc", "sciserver", "aws"])
def test__guess_host_know(host):
    # Use a new HeasarcClass object
    assert Heasarc._guess_host(host=host) == host


def test__guess_host_sciserver(monkeypatch):
    monkeypatch.setenv("HOME", "/home/idies")
    monkeypatch.setattr("os.path.exists", lambda path: path.startswith('/FTP'))
    assert Heasarc._guess_host(host=None) == 'sciserver'


@pytest.mark.parametrize("var", ["AWS_REGION", "AWS_REGION_DEFAULT", "AWS_ROLE_ARN"])
def test__guess_host_aws(monkeypatch, var):
    monkeypatch.setenv("AWS_REGION", var)
    assert Heasarc._guess_host(host=None) == 'aws'


def test_download_data__empty():
    with pytest.raises(ValueError, match="Input links table is empty"):
        Heasarc.download_data(Table())


def test_download_data__wronghost():
    with pytest.raises(
        ValueError, match="host has to be one of heasarc, sciserver, aws"
    ):
        Heasarc.download_data(Table({"id": [1]}), host="nohost")


@pytest.mark.parametrize("host", ["heasarc", "sciserver", "aws"])
def test_download_data__missingcolumn(host):
    host_col = "access_url" if host == "heasarc" else host
    with pytest.raises(
        ValueError,
        match=f"No {host_col} column found in the table. Call `~locate_data` first"
    ):
        Heasarc.download_data(Table({"id": [1]}), host=host)


def test_download_data__sciserver():
    with tempfile.TemporaryDirectory() as tmpdir:
        datadir = f'{tmpdir}/data'
        downloaddir = f'{tmpdir}/download'
        os.makedirs(datadir, exist_ok=True)
        with open(f'{datadir}/file.txt', 'w') as fp:
            fp.write('data')
        # include both a file and a directory
        tab = Table({'sciserver': [f'{tmpdir}/data/file.txt', f'{tmpdir}/data']})
        # The patch is to avoid the test that we are on sciserver
        with patch('os.path.exists') as exists:
            exists.return_value = True
            Heasarc.download_data(tab, host="sciserver", location=downloaddir)
        assert os.path.exists(f'{downloaddir}/file.txt')
        assert os.path.exists(f'{downloaddir}/data')
        assert os.path.exists(f'{downloaddir}/data/file.txt')


def test_download_data__outside_sciserver():
    with pytest.raises(
        FileNotFoundError,
        match="No data archive found. This should be run on Sciserver",
    ):
        Heasarc.download_data(
            Table({"sciserver": ["some-link"]}), host="sciserver"
        )


def test_download_data__table_row():
    with tempfile.TemporaryDirectory() as tmpdir:
        datadir = f'{tmpdir}/data'
        downloaddir = f'{tmpdir}/download'
        os.makedirs(datadir, exist_ok=True)
        with open(f'{datadir}/file.txt', 'w') as fp:
            fp.write('data')
        # include both a file and a directory
        tab = Table({'sciserver': [f'{tmpdir}/data/file.txt', f'{tmpdir}/data']})
        # The patch is to avoid the test that we are on sciserver
        with patch('os.path.exists') as exists:
            exists.return_value = True
            Heasarc.download_data(tab[0], host="sciserver", location=downloaddir)
            Heasarc.download_data(tab[1], host="sciserver", location=downloaddir)
        assert os.path.exists(f'{downloaddir}/file.txt')
        assert os.path.exists(f'{downloaddir}/data')
        assert os.path.exists(f'{downloaddir}/data/file.txt')


def test_download_data__exclude_rows_with_errors():
    with tempfile.TemporaryDirectory() as tmpdir:
        datadir = f'{tmpdir}/data'
        downloaddir = f'{tmpdir}/download'
        os.makedirs(datadir, exist_ok=True)
        with open(f'{datadir}/file.txt', 'w') as fp:
            fp.write('data')
        # include both a file and a directory
        tab = Table({
            'sciserver': [f'{tmpdir}/data/file.txt', f'{tmpdir}/data'],
            'error_message': ['', 'Error']
        })
        # The patch is to avoid the test that we are on sciserver
        with patch('os.path.exists') as exists:
            exists.return_value = True
            Heasarc.download_data(tab, host="sciserver", location=downloaddir)
        assert os.path.exists(f'{downloaddir}/file.txt')
        # data/ should be excluded because it has an error
        assert not os.path.exists(f'{downloaddir}/data')
        assert not os.path.exists(f'{downloaddir}/data/file.txt')


# S3 mock tests
s3_bucket = "nasa-heasarc"
s3_key1 = "some/location/file1.txt"
s3_key2 = "some/location/sub/file2.txt"
s3_key3 = "some/location/sub/sub2/file3.txt"
s3_dir = "some/location"


@pytest.fixture(name="s3_mock")
def _s3_mock(mocker):
    with mock_aws():
        conn = boto3.resource("s3", region_name="us-east-1")
        conn.create_bucket(Bucket=s3_bucket)
        s3_client = conn.meta.client
        s3_client.put_object(Bucket=s3_bucket, Key=s3_key1,
                             Body="my content")
        s3_client.put_object(Bucket=s3_bucket, Key=s3_key2,
                             Body="my other content")
        s3_client.put_object(Bucket=s3_bucket, Key=s3_key3,
                             Body="my other content")
        yield conn


@pytest.mark.filterwarnings("ignore::DeprecationWarning")
@pytest.mark.skipif("not DO_AWS_S3")
def test_s3_mock_basic(s3_mock):
    body = s3_mock.Object(s3_bucket, s3_key1).get()["Body"]
    content = body.read().decode("utf-8")
    assert content == "my content"


@pytest.mark.filterwarnings("ignore::DeprecationWarning")
@pytest.mark.skipif("not DO_AWS_S3")
def test_s3_mock_file(s3_mock):
    links = Table({"aws": [f"s3://{s3_bucket}/{s3_key1}"]})
    Heasarc.enable_cloud(profile=False)
    with tempfile.TemporaryDirectory() as tmpdir:
        Heasarc.download_data(links, host="aws", location=tmpdir)
        file = s3_key1.split("/")[-1]
        assert os.path.exists(f'{tmpdir}/{file}')


@pytest.mark.filterwarnings("ignore::DeprecationWarning")
@pytest.mark.skipif("not DO_AWS_S3")
def test_s3_mock_directory(s3_mock):
    links = Table({"aws": [f"s3://{s3_bucket}/{s3_dir}"]})
    Heasarc.enable_cloud(profile=False)
    with tempfile.TemporaryDirectory() as tmpdir:
        Heasarc.download_data(links, host="aws", location=tmpdir)
        assert os.path.exists(f"{tmpdir}/location")
        assert os.path.exists(f"{tmpdir}/location/file1.txt")
        assert os.path.exists(f"{tmpdir}/location/sub/file2.txt")
        assert os.path.exists(f"{tmpdir}/location/sub/sub2/file3.txt")


def test_row_count(mock_tap, mock_default_cols):
    cat = "name-1"
    assert Heasarc.count_rows(cat) == 3055
