"""Tyko client for API communication."""

import os
from types import TracebackType
from typing import Optional

import requests

from .utils import generate_random_run_name


class Run:
    """A run represents a single execution of an experiment.

    Runs are used to log parameters, metrics, and artifacts for a specific
    experiment execution. They are typically used as context managers to
    automatically handle the run lifecycle.

    Attributes:
        name: The name of this run.
        experiment: The experiment this run belongs to.

    Example:
        >>> experiment = client.create_experiment("my-experiment")
        >>> with experiment.start_run() as run:
        ...     print(run.name)  # e.g., "brave-golden-eagle"
        ...     run.log_params({"learning_rate": 0.001})
        ...     run.log_metric("accuracy", 0.95)
    """
    _experiment: "Experiment | None"

    def __init__(self, experiment: "Experiment | None" = None, name: str | None = None):
        """Initialize a run.

        Args:
            experiment: The experiment this run belongs to.
            name: Optional name for the run. If not provided, a random name
                will be generated (e.g., "brave-golden-eagle").
        """
        self._experiment = experiment
        self.name = name if name is not None else generate_random_run_name()

    def __repr__(self) -> str:
        return f"<Run name={self.name} experiment={self._experiment.name if self._experiment else 'None'}>"

    def __enter__(self) -> "Run":
        """Enter the context manager and start the run.

        Returns:
            The run instance.
        """
        return self

    def __exit__(
            self,
            exc_type: type[BaseException] | None,
            exc_value: BaseException | None,
            traceback: TracebackType | None
        ) -> bool:
        """Exit the context manager and finalize the run.

        Args:
            exc_type: The type of exception that occurred, if any.
            exc_value: The exception instance that occurred, if any.
            traceback: The traceback object for the exception, if any.

        Returns:
            False to propagate exceptions, True to suppress them.
        """
        return True


class Experiment:
    """An experiment groups related runs for comparison and organization.

    Experiments are created through the Tyko client and serve as containers
    for multiple runs. Each experiment has a unique name and can track
    multiple training runs with different parameters and configurations.

    Attributes:
        name: The name of the experiment.

    Example:
        >>> client = TykoClient()
        >>> experiment = client.create_experiment("mnist-classifier")
        >>> with experiment.start_run() as run:
        ...     run.log_params({"learning_rate": 0.001})
    """
    _client: "TykoClient | None"

    def __init__(self, name: str, client: "TykoClient | None" = None):
        """Initialize an experiment.

        Args:
            name: The name of the experiment.
            client: The Tyko client instance that created this experiment.
        """
        self.name = name
        self._client = client

    def __repr__(self) -> str:
        return f"<Experiment name={self.name}>"

    def start_run(self, name: str | None = None) -> Run:
        """Start a new run within this experiment.

        Creates and returns a new run instance that can be used to log
        parameters, metrics, and artifacts. Best used as a context manager
        to ensure proper cleanup.

        Args:
            name: Optional name for the run. If not provided, a random name
                will be generated (e.g., "brave-golden-eagle").

        Returns:
            A new Run instance associated with this experiment.

        Example:
            >>> # With auto-generated name
            >>> with experiment.start_run() as run:
            ...     print(run.name)  # e.g., "clever-mighty-wolf"
            ...     run.log_params({"batch_size": 32})
            >>>
            >>> # With custom name
            >>> with experiment.start_run(name="baseline-v1") as run:
            ...     run.log_metric("loss", 0.15)
        """
        return Run(experiment=self, name=name)



class TykoClient:
    """Client for interacting with the Tyko Labs API.

    The TykoClient is the main entry point for tracking experiments with Tyko.
    It handles authentication, experiment creation, and communication with the
    Tyko API server.

    Attributes:
        api_url: The base URL for the Tyko API.
        api_key: The API key for authentication (if provided).
        session: The requests session for API communication.

    Example:
        >>> # With API key
        >>> client = TykoClient(api_key="your-api-key")
        >>>
        >>> # Anonymous (creates temporary account)
        >>> client = TykoClient()
        >>>
        >>> # Create and track experiments
        >>> experiment = client.create_experiment("my-experiment")
        >>> with experiment.start_run() as run:
        ...     run.log_metric("accuracy", 0.95)
    """

    api_url: str

    def __init__(self, api_url: str = "https://api.tyko-labs.com", api_key: Optional[str] = None):
        """Initialize the Tyko client.

        Args:
            api_url: The base URL for the Tyko API. Defaults to the production
                server at https://api.tyko-labs.com.
            api_key: API key for authentication. If not provided, will check
                the TYKO_API_KEY environment variable. If still not found,
                an anonymous session will be created automatically.

        Example:
            >>> # With explicit API key
            >>> client = TykoClient(api_key="your-api-key")
            >>>
            >>> # Using environment variable
            >>> import os
            >>> os.environ["TYKO_API_KEY"] = "your-api-key"
            >>> client = TykoClient()
            >>>
            >>> # Anonymous session (no API key)
            >>> client = TykoClient()
        """
        self.api_url = api_url.rstrip("/")
        self.api_key = api_key or os.environ.get("TYKO_API_KEY")
        self.session = requests.Session()

        if self.api_key:
            self.session.headers.update({"Authorization": f"Bearer {self.api_key}"})

        self.session.headers.update({"Content-Type": "application/json"})

    def create_experiment(self, name: str) -> Experiment:
        """Create a new experiment.

        Creates an experiment that serves as a container for related runs.
        Multiple runs can be started from the same experiment to compare
        different configurations or training attempts.

        Args:
            name: A descriptive name for the experiment. Should be unique
                and meaningful (e.g., "mnist-classifier", "bert-fine-tuning").

        Returns:
            An Experiment instance that can be used to start runs.

        Example:
            >>> client = TykoClient()
            >>> experiment = client.create_experiment("my-ml-model")
            >>> print(experiment.name)
            my-ml-model
            >>>
            >>> with experiment.start_run() as run:
            ...     run.log_params({"learning_rate": 0.001})
        """

        return Experiment(name=name, client=self)

