"""Callback system for training monitoring and notifications."""

import abc
import json
import logging
import time
import urllib.request
import urllib.error
from typing import Dict, Any, Optional

logger = logging.getLogger(__name__)


class Callback(abc.ABC):
    """Abstract base class for training callbacks."""

    def on_train_begin(self, logs: Optional[Dict[str, Any]] = None):
        """Called at the beginning of training."""
        pass

    def on_epoch_end(self, epoch: int, logs: Optional[Dict[str, Any]] = None):
        """Called at the end of each epoch."""
        pass

    def on_train_end(self, logs: Optional[Dict[str, Any]] = None):
        """Called at the end of training."""
        pass


class Notification(Callback):
    """Sends training notifications to messaging platforms.

    Supports Slack, Discord, and Telegram via webhooks/API.

    Args:
        service: "slack", "discord", or "telegram"
        webhook_url: Webhook URL (for Slack/Discord)
        token: Bot token (for Telegram)
        chat_id: Chat ID (for Telegram)
        name: Optional name for the run to identify it in messages
    """

    def __init__(
        self,
        service: str,
        webhook_url: Optional[str] = None,
        token: Optional[str] = None,
        chat_id: Optional[str] = None,
        name: str = "Model Training",
    ):
        self.service = service.lower()
        self.webhook_url = webhook_url
        self.token = token
        self.chat_id = chat_id
        self.name = name
        self.start_time = 0.0

        if self.service == "telegram" and (not token or not chat_id):
            raise ValueError("Telegram requires both 'token' and 'chat_id'")
        
        if self.service in ["slack", "discord"] and not webhook_url:
             raise ValueError(f"{service.capitalize()} requires 'webhook_url'")

    def _send_message(self, text: str):
        """Send message to the configured service."""
        try:
            if self.service == "discord":
                self._send_discord(text)
            elif self.service == "slack":
                self._send_slack(text)
            elif self.service == "telegram":
                self._send_telegram(text)
            else:
                logger.warning(f"Unknown notification service: {self.service}")
        except Exception as e:
            logger.error(f"Failed to send notification: {e}")

    def _send_discord(self, text: str):
        data = {"content": text}
        req = urllib.request.Request(
            self.webhook_url,
            data=json.dumps(data).encode("utf-8"),
            headers={"Content-Type": "application/json", "User-Agent": "Composennent/1.0"},
        )
        with urllib.request.urlopen(req) as response:
            if response.status not in [200, 204]:
                logger.error(f"Discord webhook failed: {response.status} {response.read()}")

    def _send_slack(self, text: str):
        data = {"text": text}
        req = urllib.request.Request(
            self.webhook_url,
            data=json.dumps(data).encode("utf-8"),
            headers={"Content-Type": "application/json"},
        )
        with urllib.request.urlopen(req) as response:
             if response.status != 200:
                logger.error(f"Slack webhook failed: {response.status} {response.read()}")

    def _send_telegram(self, text: str):
        url = f"https://api.telegram.org/bot{self.token}/sendMessage"
        data = {"chat_id": self.chat_id, "text": text}
        req = urllib.request.Request(
            url,
            data=json.dumps(data).encode("utf-8"),
            headers={"Content-Type": "application/json"},
        )
        with urllib.request.urlopen(req) as response:
            if response.status != 200:
                logger.error(f"Telegram API failed: {response.status} {response.read()}")

    def on_train_begin(self, logs=None):
        self.start_time = time.time()
        self._send_message(f"🚀 **{self.name}** started!")

    def on_epoch_end(self, epoch: int, logs: Optional[Dict[str, Any]] = None):
        logs = logs or {}
        loss_info = f"Loss: {logs.get('loss', 'N/A')}"
        if 'avg_loss' in logs:
            loss_info = f"Avg Loss: {logs['avg_loss']}"
            
        self._send_message(f"📈 **{self.name}** - Epoch {epoch+1} Complete\n{loss_info}")

    def on_train_end(self, logs: Optional[Dict[str, Any]] = None):
        duration = time.time() - self.start_time
        self._send_message(f"✅ **{self.name}** finished in {duration:.2f}s!")
