import logging
import os
import time
import uuid

import pytest
import yaml

from clarifai.client.user import User

CREATE_APP_USER_ID = os.environ["CLARIFAI_USER_ID"]
NOW = uuid.uuid4().hex[:10]
CREATE_APP_ID = f"ci_input_app_{NOW}"
CREATE_DATASET_ID = "ci_input_test_dataset"
CREATE_MODEL_ID = "ci_input_test_model"
CREATE_MODEL_ID_1 = "ci_input_test_model_1"
CSV_FILE_PATH = os.path.dirname(__file__) + "/assets/sample.csv"


def create_app():
  client = User(user_id=CREATE_APP_USER_ID)
  return client.create_app(app_id=CREATE_APP_ID, base_workflow="Empty")


@pytest.mark.requires_secrets
class Testmodeltrain:
  """Tests for model training.
  """

  @classmethod
  def setup_class(self):
    self.app = create_app()
    self.input_object = self.app.inputs()
    self.dataset = self.app.create_dataset(dataset_id=CREATE_DATASET_ID)
    self.visual_classifier_model = self.app.create_model(
        model_id=CREATE_MODEL_ID, model_type_id='visual-classifier')
    self.text_classifier_model = self.app.create_model(
        model_id=CREATE_MODEL_ID_1, model_type_id='text-classifier')

  def test_model_templates(self):
    model_types = self.app.list_trainable_model_types()
    templates = self.visual_classifier_model.list_training_templates()
    assert self.visual_classifier_model.model_type_id == 'visual-classifier'  #create model test
    assert len(model_types) == 8  #list trainable model types test
    assert len(templates) >= 11  #list training templates test

  def test_model_params(self):
    model_params = self.visual_classifier_model.get_params(
        template='MMClassification_AdvancedConfig', save_to='tests/assets/model_params.yaml')
    with open('tests/assets/model_params.yaml', 'r') as file:
      expected_data = yaml.safe_load(file)
    with open(expected_data['train_params']['custom_config'], 'r') as python_file:
      custom_config = python_file.read()

    assert expected_data['train_params'].keys() == model_params[
        'train_params'].keys()  #test yaml model params
    assert len(model_params['concepts']) == 0  #test model params concepts
    assert 'custom_config' in expected_data['train_params'].keys()  #test custom config in yaml
    assert expected_data['train_params']['template'] == 'MMClassification_AdvancedConfig'  #test template param in yaml
    assert 'image_size' in expected_data['train_params'].keys()  #test image size param in yaml
    assert 'dataset_id' in expected_data.keys()  #test dataset id param in yaml
    assert '_base_' in custom_config  #test custom config script
    assert 'max_epochs' in custom_config  #test custom config script
    # cleanup
    os.remove('tests/assets/model_params.yaml')
    os.remove(expected_data['train_params']['custom_config'])

  def test_model_train(self, caplog):
    self.dataset.upload_from_csv(
        csv_path=CSV_FILE_PATH, input_type='text', csv_type='raw', labels=True)
    dataset_version = self.dataset.create_version()
    concepts = [concept.id for concept in self.app.list_concepts()]
    self.text_classifier_model.get_params(
        template='HF_GPTNeo_125m_lora', save_to='tests/assets/model_params.yaml')
    param_info = self.text_classifier_model.get_param_info(param='tokenizer_config')
    assert param_info['param'] == 'tokenizer_config'  #test get param info
    assert len(concepts) == 2  #test data upload for training
    self.text_classifier_model.update_params(
        dataset_id=CREATE_DATASET_ID,
        concepts=concepts,
        dataset_version_id=dataset_version.version.id)
    with caplog.at_level(logging.INFO):
      model_version_id = self.text_classifier_model.train()
      assert "Model Training Started" in caplog.text  #test model training

    training_status = self.text_classifier_model.training_status(version_id=model_version_id)
    assert model_version_id  #test model version id from model.train()
    assert training_status.code  #test training_status
    # NOTE (EAGLE-4139) - Immediately deleting the app causes the training to fail because the first
    # step of training requires an app, this triggers infra alerts for non graceful exit. We need to delay the deletion
    # until we get past this step.

    time.sleep(5)

    with caplog.at_level(logging.INFO):
      self.text_classifier_model.delete_version(version_id=model_version_id)
      assert "Model Version Deleted" in caplog.text  #test model version deletion
    # cleanup
    os.remove('tests/assets/model_params.yaml')

  @classmethod
  def teardown_class(self):
    self.app.delete_dataset(dataset_id=CREATE_DATASET_ID)
    self.app.delete_model(model_id=CREATE_MODEL_ID)
    self.app.delete_model(model_id=CREATE_MODEL_ID_1)
    User(user_id=CREATE_APP_USER_ID).delete_app(app_id=CREATE_APP_ID)
