"""Tests for dataset-related SDK functions."""
import pytest
from unittest.mock import MagicMock, patch

from mldash._api import _download, fetch_datasets
from mldash._errors import MLDashAuthError, MLDashUploadError


class TestFetchDatasets:
    @patch('mldash._api.requests.get')
    def test_returns_results(self, mock_get):
        mock_resp = MagicMock()
        mock_resp.status_code = 200
        mock_resp.json.return_value = {
            'results': [
                {'id': 1, 'name': 'Dataset A', 'rows': 100, 'columns': 5},
                {'id': 2, 'name': 'Dataset B', 'rows': 200, 'columns': 3},
            ]
        }
        mock_get.return_value = mock_resp

        datasets = fetch_datasets('http://localhost:8000', 'mldash_key')
        assert len(datasets) == 2
        assert datasets[0]['name'] == 'Dataset A'

    @patch('mldash._api.requests.get')
    def test_passes_project_filter(self, mock_get):
        mock_resp = MagicMock()
        mock_resp.status_code = 200
        mock_resp.json.return_value = {'results': []}
        mock_get.return_value = mock_resp

        fetch_datasets('http://localhost:8000', 'mldash_key', project_id=5)
        assert mock_get.call_args.kwargs['params'] == {'project': '5'}

    @patch('mldash._api.requests.get')
    def test_no_project_filter(self, mock_get):
        mock_resp = MagicMock()
        mock_resp.status_code = 200
        mock_resp.json.return_value = {'results': []}
        mock_get.return_value = mock_resp

        fetch_datasets('http://localhost:8000', 'mldash_key')
        assert mock_get.call_args.kwargs['params'] == {}


class TestDownload:
    @patch('mldash._api.requests.get')
    def test_download_success(self, mock_get):
        mock_resp = MagicMock()
        mock_resp.status_code = 200
        mock_resp.text = 'col1,col2\n1,2\n3,4'
        mock_get.return_value = mock_resp

        resp = _download('/api/v1/datasets/1/', 'http://localhost:8000', 'mldash_key')
        assert resp.text == 'col1,col2\n1,2\n3,4'
        mock_get.assert_called_once_with(
            'http://localhost:8000/api/v1/datasets/1/',
            headers={'X-API-Key': 'mldash_key'},
            stream=True,
            timeout=300,
        )

    @patch('mldash._api.requests.get')
    def test_download_no_api_key(self, mock_get):
        mock_resp = MagicMock()
        mock_resp.status_code = 200
        mock_get.return_value = mock_resp

        _download('/api/v1/datasets/1/', 'http://localhost:8000', None)
        assert mock_get.call_args.kwargs['headers'] == {}

    @patch('mldash._api.requests.get')
    def test_download_404(self, mock_get):
        mock_resp = MagicMock()
        mock_resp.status_code = 404
        mock_get.return_value = mock_resp

        with pytest.raises(MLDashUploadError, match='not found'):
            _download('/api/v1/datasets/999/', 'http://localhost:8000', 'mldash_key')

    @patch('mldash._api.requests.get')
    def test_download_401(self, mock_get):
        mock_resp = MagicMock()
        mock_resp.status_code = 401
        mock_get.return_value = mock_resp

        with pytest.raises(MLDashAuthError, match='Authentication required'):
            _download('/api/v1/datasets/1/', 'http://localhost:8000', 'mldash_key')

    @patch('mldash._api.requests.get')
    def test_download_403(self, mock_get):
        mock_resp = MagicMock()
        mock_resp.status_code = 403
        mock_get.return_value = mock_resp

        with pytest.raises(MLDashAuthError, match='Not authorized'):
            _download('/api/v1/datasets/1/', 'http://localhost:8000', 'mldash_key')

    @patch('mldash._api.requests.get')
    def test_download_connection_error(self, mock_get):
        import requests as req
        mock_get.side_effect = req.ConnectionError('refused')

        with pytest.raises(MLDashUploadError, match='Could not connect'):
            _download('/api/v1/datasets/1/', 'http://localhost:8000', 'mldash_key')


class TestReadDataset:
    @patch('mldash._download')
    @patch('mldash._resolve_api_key', return_value='mldash_key')
    @patch('mldash._resolve_url', return_value='http://localhost:8000')
    def test_read_dataset_returns_dataframe(self, mock_url, mock_key, mock_dl):
        mock_resp = MagicMock()
        mock_resp.text = 'a,b\n1,2\n3,4'
        mock_dl.return_value = mock_resp

        import mldash
        df = mldash.read_dataset(1)
        assert df.shape == (2, 2)
        assert list(df.columns) == ['a', 'b']
        mock_dl.assert_called_once_with(
            '/api/v1/datasets/1/',
            'http://localhost:8000', 'mldash_key',
        )

    @patch('mldash._download')
    @patch('mldash._resolve_api_key', return_value='mldash_key')
    @patch('mldash._resolve_url', return_value='http://localhost:8000')
    def test_read_dataset_public_username_slug(self, mock_url, mock_key, mock_dl):
        mock_resp = MagicMock()
        mock_resp.text = 'x,y\n5,6'
        mock_dl.return_value = mock_resp

        import mldash
        df = mldash.read_dataset('alice/her-dataset')
        mock_dl.assert_called_once_with(
            '/api/v1/datasets/public/alice/her-dataset/',
            'http://localhost:8000', 'mldash_key',
        )
        assert df.shape == (1, 2)


class TestListDatasets:
    @patch('mldash.fetch_datasets')
    @patch('mldash._resolve_api_key', return_value='mldash_key')
    @patch('mldash._resolve_url', return_value='http://localhost:8000')
    def test_list_datasets_prints_table(self, mock_url, mock_key, mock_fetch, capsys):
        mock_fetch.return_value = [
            {'id': 1, 'name': 'Sales Data', 'rows': 100, 'columns': 5,
             'dataset_type': 'training', 'visibility': 'private'},
        ]

        import mldash
        result = mldash.list_datasets()
        assert len(result) == 1
        captured = capsys.readouterr()
        assert 'Sales Data' in captured.out
        assert 'ID' in captured.out

    @patch('mldash.fetch_datasets')
    @patch('mldash._resolve_api_key', return_value='mldash_key')
    @patch('mldash._resolve_url', return_value='http://localhost:8000')
    def test_list_datasets_empty(self, mock_url, mock_key, mock_fetch, capsys):
        mock_fetch.return_value = []

        import mldash
        result = mldash.list_datasets()
        assert result == []
        captured = capsys.readouterr()
        assert 'No datasets found' in captured.out
