from __future__ import annotations

import uuid
from contextvars import ContextVar
from typing import Optional, Tuple

from ..context import get_context
from ..middleware import (
    get_call_id,
    get_request_id,
    get_sequence_number,
    set_call_id,
    set_sequence_number,
)
_TRACE_SEQUENCE: ContextVar[tuple[str, int]] = ContextVar(
    "llm_tracker_trace_sequence",
    default=("", 0),
)


def get_current_call_id() -> str:
    return get_call_id()


def set_current_call_id(call_id: str) -> None:
    set_call_id(call_id)


def _next_sequence(trace_id: str) -> int:
    current_trace, current_seq = _TRACE_SEQUENCE.get()
    if not trace_id:
        trace_id = "default-trace"

    if current_trace != trace_id:
        incoming_seq = 0
        raw_seq = get_sequence_number()
        if raw_seq:
            try:
                incoming_seq = int(raw_seq)
            except Exception:
                incoming_seq = 0
        seq = incoming_seq + 1
    else:
        seq = current_seq + 1

    _TRACE_SEQUENCE.set((trace_id, seq))
    set_sequence_number(str(seq))
    return seq


def create_call_lineage() -> Tuple[str, Optional[str], int]:
    """
    Create a new call ID with parent and per-trace sequence number.

    Parent resolution order:
    1) current call id from context
    2) request id from middleware context
    3) trace id from tracker context
    """
    ctx = get_context()
    trace_id = ctx.trace_id or get_request_id() or ""
    parent_call_id = get_current_call_id() or get_request_id() or ctx.trace_id or None
    sequence_number = _next_sequence(trace_id)
    call_id = str(uuid.uuid4())
    return call_id, parent_call_id, sequence_number
