import pandas as pd
from typing import Union
from pathlib import Path
from emissary_ml_toolkit.config.loader import ConfigLoader
from emissary_ml_toolkit.storage.manager import StorageManager
from emissary_ml_toolkit.webhooks.client import WebhookClient


class EvaluationPipeline:
    def __init__(self, config_path: Union[str, Path]):
        training_job_configs = ConfigLoader.from_yaml(config_path)
        self.storage_manager = StorageManager(training_job_configs.storage_config)
        self.webhook_client = WebhookClient(training_job_configs.project_id, training_job_configs.training_job_id)
        self.test_config = training_job_configs.test_config
        self.storage = self.storage_manager.get_storage(self.test_config.test_output.storage_ref)

    def send_testing_notification(self):
        self.webhook_client.send_training_status("Testing", False)

    def save_results(self, checkpoint: int, test_result: pd.DataFrame):
        storage_path = f"{self.test_config.test_output.base_path}/{str(checkpoint)}/test_results.csv"
        self.storage.upload_dataframe(
            df=test_result,
            storage_path=storage_path,
            index=False,
        )
        return

    def send_test_result(self, checkpoint: int, aggregated_results: dict = None, graphs: dict = None):
        self.webhook_client.send_test_result(
            checkpoint=checkpoint,
            aggregated_results=aggregated_results,
            graphs=graphs
        )
        return
