import logging
from types import TracebackType
from typing import Any, Dict, Optional, Type

from ._models import Request

logger = logging.getLogger("httpcore")


class Trace:
    def __init__(
        self,
        name: str,
        request: Optional[Request] = None,
        kwargs: Optional[Dict[str, Any]] = None,
    ) -> None:
        self.name = name
        self.trace_extension = (
            None if request is None else request.extensions.get("trace")
        )
        self.debug = logger.isEnabledFor(logging.DEBUG)
        self.kwargs = kwargs or {}
        self.return_value: Any = None
        self.should_trace = self.debug or self.trace_extension is not None

    def trace(self, name: str, info: Dict[str, Any]) -> None:
        if self.trace_extension is not None:
            self.trace_extension(name, info)

        if self.debug:
            if not info or "return_value" in info and info["return_value"] is None:
                message = name
            else:
                args = " ".join([f"{key}={value!r}" for key, value in info.items()])
                message = f"{name} {args}"
            logger.debug(message)

    def __enter__(self) -> "Trace":
        if self.should_trace:
            info = self.kwargs
            self.trace(f"{self.name}.started", info)
        return self

    def __exit__(
        self,
        exc_type: Optional[Type[BaseException]] = None,
        exc_value: Optional[BaseException] = None,
        traceback: Optional[TracebackType] = None,
    ) -> None:
        if self.should_trace:
            if exc_value is None:
                info = {"return_value": self.return_value}
                self.trace(f"{self.name}.complete", info)
            else:
                info = {"exception": exc_value}
                self.trace(f"{self.name}.failed", info)

    async def atrace(self, name: str, info: Dict[str, Any]) -> None:
        if self.trace_extension is not None:
            await self.trace_extension(name, info)

        if self.debug:
            if not info or "return_value" in info and info["return_value"] is None:
                message = name
            else:
                args = " ".join([f"{key}={value!r}" for key, value in info.items()])
                message = f"{name} {args}"
            logger.debug(message)

    async def __aenter__(self) -> "Trace":
        if self.should_trace:
            info = self.kwargs
            await self.atrace(f"{self.name}.started", info)
        return self

    async def __aexit__(
        self,
        exc_type: Optional[Type[BaseException]] = None,
        exc_value: Optional[BaseException] = None,
        traceback: Optional[TracebackType] = None,
    ) -> None:
        if self.should_trace:
            if exc_value is None:
                info = {"return_value": self.return_value}
                await self.atrace(f"{self.name}.complete", info)
            else:
                info = {"exception": exc_value}
                await self.atrace(f"{self.name}.failed", info)