"""HTTP sender for posting generations to the TrainTrack server."""

from __future__ import annotations

import os
import time
from typing import Optional, Any

import requests

from ..utils import get_logger

logger = get_logger(__name__)

# Environment variable names
TRAINTRACK_SERVER_URL = "TRAINTRACK_SERVER_URL"
TRAINTRACK_API_KEY = "TRAINTRACK_API_KEY"


class TrainTrackSender:
    """
    Minimal HTTP client for sending generations to the TrainTrack server.

    This is intentionally dumb - just POST and move on.

    Configuration can be passed directly or via environment variables:
        - TRAINTRACK_SERVER_URL: Server URL (default: http://localhost:8000)
        - TRAINTRACK_API_KEY: API key for authentication
    """

    def __init__(
        self,
        server_url: Optional[str] = None,
        api_key: Optional[str] = None,
        timeout: float = 10.0,
        max_retries: int = 2,
    ):
        """
        Initialize the TrainTrack sender.

        Args:
            server_url: URL of the TrainTrack server (or set TRAINTRACK_SERVER_URL env var)
            api_key: API key for authentication (or set TRAINTRACK_API_KEY env var)
            timeout: Request timeout in seconds
            max_retries: Number of retries on failure
        """
        # Use explicit args or fall back to environment variables
        self.server_url = (
            server_url
            or os.environ.get(TRAINTRACK_SERVER_URL)
            or "https://traintracklabs.com"
        ).rstrip("/")
        self.api_key = api_key or os.environ.get(TRAINTRACK_API_KEY)
        self.timeout = timeout
        self.max_retries = max_retries

    def send_generations(
        self,
        run_name: str,
        checkpoint_tag: str,
        step: int,
        generations: list[dict[str, Any]],
        epoch: Optional[int] = None,
        metadata: Optional[dict[str, Any]] = None,
    ) -> bool:
        """
        Send generations to the server.

        Args:
            run_name: Name of the training run
            checkpoint_tag: Unique tag for this checkpoint
            step: Current training step
            generations: List of generation dicts with format:
                {
                    "prompt": {
                        "id": str,
                        "dataset_name": str,
                        "prompt": str,
                        "category": Optional[str],
                        "reference": Optional[str],
                        "competitor_output": Optional[str],
                    },
                    "output": str,
                    "generation_time_ms": Optional[int],
                }
            epoch: Optional epoch number
            metadata: Optional metadata dict

        Returns:
            True if successful, False otherwise
        """
        payload = {
            "run_name": run_name,
            "checkpoint_tag": checkpoint_tag,
            "step": step,
            "epoch": epoch,
            "generations": generations,
            "metadata": metadata,
        }

        url = f"{self.server_url}/api/ingest"

        # Build headers with optional API key
        headers = {"Content-Type": "application/json"}
        if self.api_key:
            headers["X-API-Key"] = self.api_key

        for attempt in range(self.max_retries + 1):
            try:
                response = requests.post(
                    url,
                    json=payload,
                    headers=headers,
                    timeout=self.timeout,
                )

                if response.status_code == 200:
                    data = response.json()
                    if data.get("success"):
                        logger.info(
                            f"Sent {len(generations)} generations for step {step}"
                        )
                        return True
                    else:
                        logger.warning(f"Server returned error: {data.get('message')}")
                else:
                    logger.warning(
                        f"Server returned status {response.status_code}: {response.text[:200]}"
                    )

            except requests.exceptions.Timeout:
                logger.warning(f"Timeout sending to server (attempt {attempt + 1})")
            except requests.exceptions.ConnectionError:
                logger.warning(f"Connection error (attempt {attempt + 1})")
            except Exception as e:
                logger.error(f"Error sending generations: {e}")

            if attempt < self.max_retries:
                time.sleep(0.5 * (attempt + 1))

        logger.error(
            f"Failed to send generations after {self.max_retries + 1} attempts"
        )
        return False

    def check_health(self) -> bool:
        """Check if server is healthy."""
        try:
            response = requests.get(
                f"{self.server_url}/health",
                timeout=5.0,
            )
            return response.status_code == 200
        except Exception:
            return False

    def verify_connection(self) -> bool:
        """
        Verify connection and API key validity.

        Returns:
            True if valid

        Raises:
            ConnectionError: If server is unreachable
            ValueError: If API key is invalid
        """
        url = f"{self.server_url}/api/auth/verify-key"

        headers = {}
        if self.api_key:
            headers["X-API-Key"] = self.api_key

        try:
            response = requests.get(
                url,
                headers=headers,
                timeout=5.0,
            )

            if response.status_code == 200:
                data = response.json()
                logger.info(f"Connected to TrainTrack as user: {data['user']['email']}")
                return True

            if response.status_code == 401:
                raise ValueError("Invalid API Key. Please check your configuration.")

            raise ConnectionError(
                f"Server returned status {response.status_code}: {response.text}"
            )

        except requests.exceptions.RequestException as e:
            raise ConnectionError(f"Could not connect to TrainTrack server: {e}")
