from pathlib import Path
import responses
from requests.exceptions import RequestException
from unittest.mock import patch, create_autospec

import pytest

from anylearn.applications.quickstart import quick_evaluate
from anylearn.config import init_sdk
from anylearn.interfaces import EvaluateTask, EvaluateTaskState
from anylearn.interfaces.resource import (
    Algorithm,
    Dataset,
    Model,
    ResourceState,
    ResourceVisibility
)
from anylearn.utils.errors import (
    AnyLearnException,
    AnyLearnNotSupportedException
)
from tests.base_test_case import BaseTestCase


FAKE_ALGO_OK = Path(__file__).parent / "fake_algo_ok"
FAKE_ALGO_KO = Path(__file__).parent / "fake_algo_ko"
FAKE_DSET = Path(__file__).parent / "fake_dset"
FAKE_MODEL = Path(__file__).parent / "fake_model"


class TestQuickEvaluate(BaseTestCase):
    @pytest.mark.skip()
    @patch('anylearn.interfaces.resource.resource_uploader.ResourceUploader')
    @responses.activate
    def test_quick_evaluate_ok(self, MockResourceUploader):
        @create_autospec
        def mock_upload_ok(resource_id, chunks):
            return True

        # Mock uploader with stubbed upload function
        uploader = MockResourceUploader()
        uploader.run = mock_upload_ok

        # Stub mirror list
        responses.add(responses.GET, url=self._url("mirror/list"),
                      json=[
                          {
                              'id': "MIRR001",
                              'name': "Test001",
                              'requirements': "test==0.0.1",
                          },
                          {
                              'id': "MIRR002",
                              'name': "QUICKSTART",
                              'requirements': "test==0.0.2",
                          },
                      ],
                      status=200)
        # Stub algo creation
        responses.add(responses.POST, url=self._url("algorithm/add"),
                      json={'data': "ALGO001", 'message': "算法添加成功"},
                      status=200)
        # Stub model creation
        responses.add(responses.POST, url=self._url("model/add"),
                      json={'data': "MODE001", 'message': "模型添加成功"},
                      status=200)
        # Stub dset creation
        responses.add(responses.POST, url=self._url("dataset/add"),
                      json={'data': "DSET001", 'message': "数据集添加成功"},
                      status=200)
        # Stub algo upload finish
        responses.add(responses.POST, url=self._url("resource/upload_finish"),
                      json={
                          'msg': "上传已完成，稍后可查看上传结果"
                      },
                      status=200)
        # Stub model upload finish
        responses.add(responses.POST, url=self._url("resource/upload_finish"),
                      json={
                          'msg': "上传已完成，稍后可查看上传结果"
                      },
                      status=200)
        # Stub dset upload finish
        responses.add(responses.POST, url=self._url("resource/upload_finish"),
                      json={
                          'msg': "上传已完成，稍后可查看上传结果"
                      },
                      status=200)
        # Stub algo query
        responses.add(responses.GET, url=self._url("algorithm/query?id=ALGO001"),
                      match_querystring=True, json=[{
                          'id': "ALGO001",
                          'name': "TestAlgo1",
                          'description': "test",
                          'state': ResourceState.READY,
                          'visibility': ResourceVisibility.PUBLIC,
                          'upload_time': "2020-01-01 00:00",
                          'filename': "buzhongyao",
                          'is_zipfile': 1,
                          'file': "USER001/algos/ALGO001/buzhongyao",
                          'size': 250.41,
                          'creator_id': "USER001",
                          'node_id': "NODE001",
                          'owner': ["USER001"],
                          'tags': "tag1, tag2",
                          'train_params': '[{"test": "test"}]',
                          'evaluate_params': '[{"test": "test"}]',
                          'mirror_id': "MIRR002",
                          'follows_anylearn_norm': False,
                          'entrypoint_training': "python train.py",
                          'output_training': "output",
                          'entrypoint_evaluation': "python eval.py",
                          'output_evaluation': "output",
                      }],
                      status=200)
        # Stub model query
        responses.add(responses.GET, url=self._url("model/query?id=MODE001"),
                      match_querystring=True, json=[{
                          'id': "MODE001",
                          'name': "TestModel1",
                          'description': "test",
                          'state': ResourceState.READY,
                          'visibility': ResourceVisibility.PUBLIC,
                          'upload_time': "2020-01-01 00:00",
                          'filename': "buzhongyao",
                          'is_zipfile': 1,
                          'file': "USER001/models/MODE001/buzhongyao",
                          'size': 250.41,
                          'creator_id': "USER001",
                          'node_id': "NODE001",
                          'owner': ["USER001"],
                          'algorithm_id': "ALGO001",
                      }],
                      status=200)
        # Stub dset query
        responses.add(responses.GET, url=self._url("dataset/query?id=DSET001"),
                      match_querystring=True, json=[{
                          'id': "DSET001",
                          'name': "TestDset1",
                          'description': "test",
                          'state': ResourceState.READY,
                          'visibility': ResourceVisibility.PUBLIC,
                          'upload_time': "2020-01-01 00:00",
                          'filename': "buzhongyao",
                          'is_zipfile': 1,
                          'file': "USER001/datasets/DSET001/buzhongyao",
                          'size': 250.41,
                          'creator_id': "USER001",
                          'node_id': "NODE001",
                          'owner': ["USER001"],
                      }],
                      status=200)
        # Stub eval task creation
        responses.add(responses.POST, url=self._url("evaluate_task/add"),
                      json={'data': "EVAL001", 'message': "服务添加成功"},
                      status=200)
        # Stub eval task query
        responses.add(responses.GET, url=self._url("evaluate_task/query?id=EVAL001"),
                      match_querystring=True, json={
                          'id': "EVAL001",
                          'name': "TestEvaluateTask1",
                          'description': "test",
                          'state': EvaluateTaskState.RUNNING,
                          'visibility': 1,
                          'creator_id': "USER001",
                          'owner': ["USER001"],
                          'create_time': "2021-04-01 23:59:59",
                          'finish_time': "",
                          'params_list': '[{"model_id": "MODE001", "files": "DSET001", "evaluate_params": {"dataset": "$DSET001", "model_path": "$MODE001"}}]',
                          'secret_key': "SECRET",
                          'results': [
                              {
                                  'args': '{"dataset": "$DSET001", "model_path": "$MODE001"}',
                                  'envs': "",
                                  'files': "DSET001",
                                  'gpu_num': 1,
                                  'id': "EVRE001",
                                  'model_id': "MODE001",
                                  'results': "",
                              }
                          ],
                      },
                      status=200)

        eval_task, algo, model, dset = quick_evaluate(
            model_dir=FAKE_MODEL,
            algorithm_dir=FAKE_ALGO_OK,
            dataset_dir=FAKE_DSET,
            entrypoint="python main.py",
            output="output",
            resource_uploader=uploader,
            resource_polling=0.1,
            model_hyperparam_name="model_path",
            dataset_hyperparam_name="dataset",
            hyperparams={'test_param': "test"}
        )

        self.assertIsInstance(eval_task, EvaluateTask)
        self.assertEqual(eval_task.id, "EVAL001")
        self.assertTrue((FAKE_ALGO_OK / "anylearn_tools.py").exists())
        self.assertIsInstance(algo, Algorithm)
        self.assertEqual(algo.id, "ALGO001")
        self.assertIsInstance(model, Model)
        self.assertEqual(model.id, "MODE001")
        self.assertIsInstance(dset, Dataset)
        self.assertEqual(dset.id, "DSET001")

    @pytest.mark.skip()
    def test_quick_evaluate_ko_algo_missing_requirements_txt(self):
        with self.assertRaises(AnyLearnException) as ctx:
            eval_task = quick_evaluate(model_dir=FAKE_MODEL,
                                       algorithm_dir=FAKE_ALGO_KO,
                                       dataset_dir=FAKE_DSET,
                                       entrypoint="python main.py",
                                       output="output")
        e = ctx.exception
        self.assertIsInstance(e, AnyLearnException)
        self.assertEqual(e.msg, "Missing 'requirements.txt' in algorithm folder")

    @pytest.mark.skip()
    @responses.activate
    def test_quick_evaluate_ko_mirror_not_found(self):
        responses.add(responses.GET, url=self._url("mirror/list"),
                      json=[
                          {
                              'id': "MIRR001",
                              'name': "Test001",
                              'requirements': "test==0.0.1",
                          },
                          {
                              'id': "MIRR002",
                              'name': "Test002",
                              'requirements': "test==0.0.2",
                          },
                      ],
                      status=200)
        with self.assertRaises(AnyLearnNotSupportedException) as ctx:
            eval_task = quick_evaluate(model_dir=FAKE_MODEL,
                                       algorithm_dir=FAKE_ALGO_OK,
                                       dataset_dir=FAKE_DSET,
                                       entrypoint="python main.py",
                                       output="output")
        e = ctx.exception
        self.assertIsInstance(e, AnyLearnNotSupportedException)
        self.assertEqual(e.msg, "Container for `QUICKSTART` is not supported by the connected backend.")

    @pytest.mark.skip()
    @patch('anylearn.interfaces.resource.resource_uploader.ResourceUploader')
    @responses.activate
    def test_quick_evaluate_ko_algo_on_error(self, MockResourceUploader):
        @create_autospec
        def mock_upload_ok(resource_id, chunks):
            return True

        # Mock uploader with stubbed upload function
        uploader = MockResourceUploader()
        uploader.run = mock_upload_ok

        # Stub mirror list
        responses.add(responses.GET, url=self._url("mirror/list"),
                      json=[
                          {
                              'id': "MIRR001",
                              'name': "Test001",
                              'requirements': "test==0.0.1",
                          },
                          {
                              'id': "MIRR002",
                              'name': "QUICKSTART",
                              'requirements': "test==0.0.2",
                          },
                      ],
                      status=200)
        # Stub algo creation
        responses.add(responses.POST, url=self._url("algorithm/add"),
                      json={'data': "ALGO250", 'message': "算法添加成功"},
                      status=200)
        # Stub model creation
        responses.add(responses.POST, url=self._url("model/add"),
                      json={'data': "MODE001", 'message': "模型添加成功"},
                      status=200)
        # Stub dset creation
        responses.add(responses.POST, url=self._url("dataset/add"),
                      json={'data': "DSET001", 'message': "数据集添加成功"},
                      status=200)
        # Stub algo upload finish
        responses.add(responses.POST, url=self._url("resource/upload_finish"),
                      json={
                          'msg': "上传已完成，稍后可查看上传结果"
                      },
                      status=200)
        # Stub model upload finish
        responses.add(responses.POST, url=self._url("resource/upload_finish"),
                      json={
                          'msg': "上传已完成，稍后可查看上传结果"
                      },
                      status=200)
        # Stub dset upload finish
        responses.add(responses.POST, url=self._url("resource/upload_finish"),
                      json={
                          'msg': "上传已完成，稍后可查看上传结果"
                      },
                      status=200)
        # Stub algo query
        responses.add(responses.GET, url=self._url("algorithm/query?id=ALGO250"),
                      match_querystring=True, json=[{
                          'id': "ALGO250",
                          'name': "TestAlgo1",
                          'description': "test",
                          'state': ResourceState.ERROR,
                          'visibility': ResourceVisibility.PUBLIC,
                          'upload_time': "2020-01-01 00:00",
                          'filename': "buzhongyao",
                          'is_zipfile': 1,
                          'file': "USER001/algos/ALGO250/buzhongyao",
                          'size': 250.41,
                          'creator_id': "USER001",
                          'node_id': "NODE001",
                          'owner': ["USER001"],
                          'tags': "tag1, tag2",
                          'train_params': '[{"test": "test"}]',
                          'evaluate_params': '[{"test": "test"}]',
                          'mirror_id': "MIRR002",
                          'follows_anylearn_norm': False,
                          'entrypoint_training': "python train.py",
                          'output_training': "output",
                          'entrypoint_evaluation': "python eval.py",
                          'output_evaluation': "output",
                      }],
                      status=200)

        with self.assertRaises(AnyLearnException) as ctx:
            eval_task = quick_evaluate(model_dir=FAKE_MODEL,
                                       algorithm_dir=FAKE_ALGO_OK,
                                       dataset_dir=FAKE_DSET,
                                       entrypoint="python main.py",
                                       output="output",
                                       resource_uploader=uploader,
                                       resource_polling=0.1)
        e = ctx.exception
        self.assertIsInstance(e, AnyLearnException)
        self.assertEqual(e.msg, "Error occured when uploading algorithm")

    @pytest.mark.skip()
    @patch('anylearn.interfaces.resource.resource_uploader.ResourceUploader')
    @responses.activate
    def test_quick_evaluate_ko_dset_on_error(self, MockResourceUploader):
        @create_autospec
        def mock_upload_ok(resource_id, chunks):
            return True

        # Mock uploader with stubbed upload function
        uploader = MockResourceUploader()
        uploader.run = mock_upload_ok

        # Stub mirror list
        responses.add(responses.GET, url=self._url("mirror/list"),
                      json=[
                          {
                              'id': "MIRR001",
                              'name': "Test001",
                              'requirements': "test==0.0.1",
                          },
                          {
                              'id': "MIRR002",
                              'name': "QUICKSTART",
                              'requirements': "test==0.0.2",
                          },
                      ],
                      status=200)
        # Stub algo creation
        responses.add(responses.POST, url=self._url("algorithm/add"),
                      json={'data': "ALGO001", 'message': "算法添加成功"},
                      status=200)
        # Stub model creation
        responses.add(responses.POST, url=self._url("model/add"),
                      json={'data': "MODE001", 'message': "模型添加成功"},
                      status=200)
        # Stub dset creation
        responses.add(responses.POST, url=self._url("dataset/add"),
                      json={'data': "DSET250", 'message': "数据集添加成功"},
                      status=200)
        # Stub algo upload finish
        responses.add(responses.POST, url=self._url("resource/upload_finish"),
                      json={
                          'msg': "上传已完成，稍后可查看上传结果"
                      },
                      status=200)
        # Stub model upload finish
        responses.add(responses.POST, url=self._url("resource/upload_finish"),
                      json={
                          'msg': "上传已完成，稍后可查看上传结果"
                      },
                      status=200)
        # Stub dset upload finish
        responses.add(responses.POST, url=self._url("resource/upload_finish"),
                      json={
                          'msg': "上传已完成，稍后可查看上传结果"
                      },
                      status=200)
        # Stub model query
        responses.add(responses.GET, url=self._url("model/query?id=MODE001"),
                      match_querystring=True, json=[{
                          'id': "MODE001",
                          'name': "TestModel1",
                          'description': "test",
                          'state': ResourceState.READY,
                          'visibility': ResourceVisibility.PUBLIC,
                          'upload_time': "2020-01-01 00:00",
                          'filename': "buzhongyao",
                          'is_zipfile': 1,
                          'file': "USER001/models/MODE001/buzhongyao",
                          'size': 250.41,
                          'creator_id': "USER001",
                          'node_id': "NODE001",
                          'owner': ["USER001"],
                          'algorithm_id': "ALGO001",
                      }],
                      status=200)
        # Stub algo query
        responses.add(responses.GET, url=self._url("algorithm/query?id=ALGO001"),
                      match_querystring=True, json=[{
                          'id': "ALGO001",
                          'name': "TestAlgo1",
                          'description': "test",
                          'state': ResourceState.READY,
                          'visibility': ResourceVisibility.PUBLIC,
                          'upload_time': "2020-01-01 00:00",
                          'filename': "buzhongyao",
                          'is_zipfile': 1,
                          'file': "USER001/algos/ALGO001/buzhongyao",
                          'size': 250.41,
                          'creator_id': "USER001",
                          'node_id': "NODE001",
                          'owner': ["USER001"],
                          'tags': "tag1, tag2",
                          'train_params': '[{"test": "test"}]',
                          'evaluate_params': '[{"test": "test"}]',
                          'mirror_id': "MIRR002",
                          'follows_anylearn_norm': False,
                          'entrypoint_training': "python train.py",
                          'output_training': "output",
                          'entrypoint_evaluation': "python eval.py",
                          'output_evaluation': "output",
                      }],
                      status=200)
        # Stub dset query
        responses.add(responses.GET, url=self._url("dataset/query?id=DSET250"),
                      match_querystring=True, json=[{
                          'id': "DSET250",
                          'name': "TestDset1",
                          'description': "test",
                          'state': ResourceState.ERROR,
                          'visibility': ResourceVisibility.PUBLIC,
                          'upload_time': "2020-01-01 00:00",
                          'filename': "buzhongyao",
                          'is_zipfile': 1,
                          'file': "USER001/datasets/DSET250/buzhongyao",
                          'size': 250.41,
                          'creator_id': "USER001",
                          'node_id': "NODE001",
                          'owner': ["USER001"],
                      }],
                      status=200)

        with self.assertRaises(AnyLearnException) as ctx:
            eval_task = quick_evaluate(model_dir=FAKE_MODEL,
                                       algorithm_dir=FAKE_ALGO_OK,
                                       dataset_dir=FAKE_DSET,
                                       entrypoint="python main.py",
                                       output="output",
                                       resource_uploader=uploader,
                                       resource_polling=0.1)
        e = ctx.exception
        self.assertIsInstance(e, AnyLearnException)
        self.assertEqual(e.msg, "Error occured when uploading dataset")

    @pytest.mark.skip()
    @patch('anylearn.interfaces.resource.resource_uploader.ResourceUploader')
    @responses.activate
    def test_quick_evaluate_ko_model_on_error(self, MockResourceUploader):
        @create_autospec
        def mock_upload_ok(resource_id, chunks):
            return True

        # Mock uploader with stubbed upload function
        uploader = MockResourceUploader()
        uploader.run = mock_upload_ok

        # Stub mirror list
        responses.add(responses.GET, url=self._url("mirror/list"),
                      json=[
                          {
                              'id': "MIRR001",
                              'name': "Test001",
                              'requirements': "test==0.0.1",
                          },
                          {
                              'id': "MIRR002",
                              'name': "QUICKSTART",
                              'requirements': "test==0.0.2",
                          },
                      ],
                      status=200)
        # Stub algo creation
        responses.add(responses.POST, url=self._url("algorithm/add"),
                      json={'data': "ALGO001", 'message': "算法添加成功"},
                      status=200)
        # Stub model creation
        responses.add(responses.POST, url=self._url("model/add"),
                      json={'data': "MODE250", 'message': "模型添加成功"},
                      status=200)
        # Stub dset creation
        responses.add(responses.POST, url=self._url("dataset/add"),
                      json={'data': "DSET001", 'message': "数据集添加成功"},
                      status=200)
        # Stub algo upload finish
        responses.add(responses.POST, url=self._url("resource/upload_finish"),
                      json={
                          'msg': "上传已完成，稍后可查看上传结果"
                      },
                      status=200)
        # Stub model upload finish
        responses.add(responses.POST, url=self._url("resource/upload_finish"),
                      json={
                          'msg': "上传已完成，稍后可查看上传结果"
                      },
                      status=200)
        # Stub dset upload finish
        responses.add(responses.POST, url=self._url("resource/upload_finish"),
                      json={
                          'msg': "上传已完成，稍后可查看上传结果"
                      },
                      status=200)
        # Stub algo query
        responses.add(responses.GET, url=self._url("algorithm/query?id=ALGO001"),
                      match_querystring=True, json=[{
                          'id': "ALGO001",
                          'name': "TestAlgo1",
                          'description': "test",
                          'state': ResourceState.READY,
                          'visibility': ResourceVisibility.PUBLIC,
                          'upload_time': "2020-01-01 00:00",
                          'filename': "buzhongyao",
                          'is_zipfile': 1,
                          'file': "USER001/algos/ALGO001/buzhongyao",
                          'size': 250.41,
                          'creator_id': "USER001",
                          'node_id': "NODE001",
                          'owner': ["USER001"],
                          'tags': "tag1, tag2",
                          'train_params': '[{"test": "test"}]',
                          'evaluate_params': '[{"test": "test"}]',
                          'mirror_id': "MIRR002",
                          'follows_anylearn_norm': False,
                          'entrypoint_training': "python train.py",
                          'output_training': "output",
                          'entrypoint_evaluation': "python eval.py",
                          'output_evaluation': "output",
                      }],
                      status=200)
        # Stub model query
        responses.add(responses.GET, url=self._url("model/query?id=MODE250"),
                      match_querystring=True, json=[{
                          'id': "MODE250",
                          'name': "TestModel1",
                          'description': "test",
                          'state': ResourceState.ERROR,
                          'visibility': ResourceVisibility.PUBLIC,
                          'upload_time': "2020-01-01 00:00",
                          'filename': "buzhongyao",
                          'is_zipfile': 1,
                          'file': "USER001/models/MODE250/buzhongyao",
                          'size': 250.41,
                          'creator_id': "USER001",
                          'node_id': "NODE001",
                          'owner': ["USER001"],
                          'algorithm_id': "ALGO001",
                      }],
                      status=200)

        with self.assertRaises(AnyLearnException) as ctx:
            eval_task = quick_evaluate(model_dir=FAKE_MODEL,
                                       algorithm_dir=FAKE_ALGO_OK,
                                       dataset_dir=FAKE_DSET,
                                       entrypoint="python main.py",
                                       output="output",
                                       resource_uploader=uploader,
                                       resource_polling=0.1)
        e = ctx.exception
        self.assertIsInstance(e, AnyLearnException)
        self.assertEqual(e.msg, "Error occured when uploading model")
