"""
contextvars를 사용한 Trace Context 관리자
FastAPI의 비동기 작업들 간에 trace context를 관리함

역할: HTTP 요청부터 API 호출, Bedrock 호출까지 동일한 trace_id를 유지하여
      분산 추적을 가능하게 함
"""

from contextvars import ContextVar
from typing import Optional
from ..types.interfaces import TraceContext
from .id_generator import generate_trace_id, generate_span_id

# ContextVar를 사용하여 비동기 컨텍스트에서 trace context 관리
_trace_context: ContextVar[Optional[TraceContext]] = ContextVar('trace_context', default=None)


def get_trace_context() -> Optional[TraceContext]:
    """
    현재 trace context를 조회함

    역할: 현재 실행 중인 비동기 작업의 trace context를 가져옴
    사용: API 호출이나 Bedrock 호출 시 parent span을 찾기 위해 사용됨

    Returns:
        현재 trace context 또는 추적 중이 아니면 None
    """
    return _trace_context.get()


def set_trace_context(context: Optional[TraceContext]) -> None:
    """
    trace context를 설정함

    역할: 특정 trace context를 현재 실행 컨텍스트에 설정함
    사용: HTTP 요청 처리 시작 시 root context를 설정함

    Args:
        context: 설정할 Trace context
    """
    _trace_context.set(context)


def create_root_context() -> TraceContext:
    """
    새로운 root trace context를 생성함 (HTTP 요청용)

    역할: HTTP 요청이 시작될 때 새로운 추적 세션을 시작함
    사용: HTTP Middleware에서 각 HTTP 요청마다 호출됨

    Returns:
        생성된 trace_id와 span_id를 가진 새 trace context
    """
    return TraceContext(
        trace_id=generate_trace_id(),
        span_id=generate_span_id(),
        parent_span_id=None
    )


def create_child_context() -> Optional[TraceContext]:
    """
    현재 context로부터 child span context를 생성함

    역할: API 호출이나 Bedrock 호출 시 child span을 생성함
    사용: HttpClientCollector, BedrockCollector에서 호출되어
          parent span(HTTP 요청)과 연결된 child span을 만듦

    Returns:
        같은 trace_id를 가지지만 새 span_id를 가진 trace context
        또는 현재 context가 없으면 None
    """
    current_context = get_trace_context()

    if not current_context:
        return None

    return TraceContext(
        trace_id=current_context.trace_id,
        span_id=generate_span_id(),
        parent_span_id=current_context.span_id
    )


def get_current_trace_id() -> Optional[str]:
    """
    현재 trace ID를 조회함

    역할: 현재 실행 중인 요청의 trace ID를 반환함
    사용: 로깅 시 trace ID를 함께 기록하고 싶을 때 사용됨

    Returns:
        현재 trace ID 또는 None
    """
    context = get_trace_context()
    return context.trace_id if context else None


def get_current_span_id() -> Optional[str]:
    """
    현재 span ID를 조회함

    역할: 현재 실행 중인 작업의 span ID를 반환함

    Returns:
        현재 span ID 또는 None
    """
    context = get_trace_context()
    return context.span_id if context else None


def get_current_parent_span_id() -> Optional[str]:
    """
    현재 parent span ID를 조회함

    역할: 현재 span의 부모 span ID를 반환함

    Returns:
        현재 parent span ID 또는 None
    """
    context = get_trace_context()
    return context.parent_span_id if context else None


def is_in_traced_context() -> bool:
    """
    현재 추적 context 내에 있는지 확인함

    역할: 현재 HTTP 요청 처리 중인지 확인함

    Returns:
        추적 중이면 True, 아니면 False
    """
    return get_trace_context() is not None
