import os
import time
import requests
from emissary_ml_toolkit.webhooks.protocol import WebhookEvent, WebhookPayload, WebhookEventType, TrainingProgress
from emissary_ml_toolkit.config.env import load_env_file

load_env_file()


class WebhookClient:
    def __init__(self,
                 project_id: str,
                 training_job_id: str,
                 endpoint_url: str = None,
                 api_key: str = None,
                 max_retries: int = 3,
                 timeout: int = 10,):
        if not endpoint_url:
            endpoint_url = os.environ.get("WEBHOOK_ENDPOINT_URL")
        if not api_key:
            api_key = os.environ.get("WEBHOOK_API_KEY", None)

        self.project_id = project_id
        self.training_job_id = training_job_id
        self.endpoint_url = endpoint_url
        self.headers = {
            "Content-Type": "application/json",
            "X-Api-Key": api_key
        }
        self.max_retries = max_retries
        self.timeout = timeout

    def _send_webhook(self, event: WebhookEvent):
        payload = WebhookPayload(
            service_id=self.project_id,
            training_job_id=self.training_job_id,
            event_type=event.event_type,
            event=event,
        )
        for attempt in range(self.max_retries):
            try:
                response = requests.post(
                    self.endpoint_url,
                    data=payload.model_dump_json(),
                    headers=self.headers,
                    timeout=self.timeout,
                )
                if response.status_code == 200:
                    return
                if response.status_code >= 500:
                    time.sleep(2 ** attempt)
                    continue
                raise ValueError(response.text)
            except requests.exceptions.RequestException as e:
                if attempt < self.max_retries - 1:
                    time.sleep(2 ** attempt)
        return

    def send_dataset_upload(self, dataset_id: str):
        event = WebhookEvent(
            event_type=WebhookEventType.DATASET_UPLOADED,
            dataset_id=dataset_id,
        )
        self._send_webhook(event)

    def send_checkpoint_upload(self, checkpoint: int):
        event = WebhookEvent(
            event_type=WebhookEventType.CHECKPOINT_UPLOADED,
            training_id=self.training_job_id,
            checkpoint=checkpoint,
        )
        self._send_webhook(event)

    def send_training_progress(self, training_progress: TrainingProgress, training_loss: float):
        event = WebhookEvent(
            event_type=WebhookEventType.TRAINING_PROGRESS_AND_LOSS_UPDATE,
            training_id=self.training_job_id,
            training_progress=training_progress,
            training_loss=training_loss,
        )
        self._send_webhook(event)

    def send_training_status(self, status: str, kill_instance: bool = False):
        event = WebhookEvent(
            event_type=WebhookEventType.TRAINING_STATUS_UPDATE,
            training_id=self.training_job_id,
            status=status,
            kill_instance=kill_instance,
        )
        self._send_webhook(event)

    def send_test_result(self, checkpoint: int, aggregated_results: dict = None, graphs: dict = None):
        event = WebhookEvent(
            event_type=WebhookEventType.MODEL_TEST_RESULTS,
            training_id=self.training_job_id,
            checkpoint=checkpoint,
            test_results=aggregated_results,    # later change to aggregated_results in platform and library
            graphs=graphs,
        )
        self._send_webhook(event)


    def send_webhook_event(self):
        pass