import os
import time

from scale_gp import SGPClient

client = SGPClient(api_key=api_key)

autogenerated_evaluation_dataset = client.evaluation_datasets.create(
    account_id=account_id,
    name="autogenerated_eval_dataset",
    schema_type="GENERATION",
    type="autogenerated",
    knowledge_base_id=knowledge_base_id,
)
print(autogenerated_evaluation_dataset)

generation_job = client.evaluation_datasets.generation_jobs.create(
    evaluation_dataset_id=autogenerated_evaluation_dataset.id,
    num_test_cases=3,
    group_by_artifact_id=False,
)
while True:
    generation_job = client.evaluation_datasets.generation_jobs.retrieve(
        generation_job_id=generation_job.generation_job_id,
        evaluation_dataset_id=autogenerated_evaluation_dataset.id
    )
    if generation_job.status == "Pending":
        print("generating test cases...")
        time.sleep(5)
    else:
        break
print(generation_job)
# view autogenerated test cases
test_cases = client.evaluation_datasets.autogenerated_draft_test_cases.list(
    evaluation_dataset_id=autogenerated_evaluation_dataset.id
)
print(test_cases.items)

for test_case in test_cases.items:
    client.evaluation_datasets.autogenerated_draft_test_cases.approve(
        evaluation_dataset_id=autogenerated_evaluation_dataset.id,
        autogenerated_draft_test_case_id=test_case.id,
    )

published_dataset_response = client.evaluation_datasets.publish(
    evaluation_dataset_id=autogenerated_evaluation_dataset.id,
)
print(published_dataset_response)

EvaluationDataset(
    id='da20895d-1784-4d40-a403-e7b5f8a2ce78',
    account_id='66049ada2fc77c99ef015be7',
    created_at=datetime.datetime(2024, 9, 26, 20, 6, 46, 689131),
    created_by_user_id='42a5c8af-f698-43d0-923e-ba70102a2887',
    draft=None,
    name='autogenerated_eval_dataset',
    schema_type='GENERATION',
    updated_at=datetime.datetime(2024, 9, 26, 20, 6, 46, 689131),
    archived_at=None,
    evaluation_dataset_metadata={
        'knowledge_base_name': 'Zelda OOT Walkthroughs'
    },
    knowledge_base_id='78447679-4cbf-4c27-af40-d4f1ec9f6911',
    out_of_date=None,
    vendor=None
)