from __future__ import annotations

from contextlib import contextmanager
import json
import os
from typing import Any, Iterator, Mapping, Optional

from dotenv import load_dotenv


load_dotenv()

try:
    from langfuse import propagate_attributes
except Exception:  # pragma: no cover - optional dependency fallback
    propagate_attributes = None  # type: ignore[assignment]


def _langfuse_is_configured() -> bool:
    return bool(os.getenv("LANGFUSE_PUBLIC_KEY") and os.getenv("LANGFUSE_SECRET_KEY"))


def _stringify_metadata(
    metadata: Optional[Mapping[str, Any]],
) -> Optional[dict[str, str]]:
    if not metadata:
        return None

    result: dict[str, str] = {}
    for key, value in metadata.items():
        if value is None:
            continue
        if isinstance(value, str):
            result[key] = value
            continue
        try:
            result[key] = json.dumps(value, ensure_ascii=False)
        except Exception:
            result[key] = str(value)

    return result or None


@contextmanager
def propagate_conversation_session(
    *,
    conversation_id: str,
    metadata: Optional[Mapping[str, Any]] = None,
    tags: Optional[list[str]] = None,
) -> Iterator[None]:
    """Attach `session_id=conversation_id` to existing spans in the current execution scope.

    SimpleLLMFunc already creates the turn/message spans for us. We only need to make sure
    all spans emitted while handling the same conversation inherit the same Langfuse session.
    """

    if not _langfuse_is_configured() or propagate_attributes is None:
        yield
        return

    propagated_metadata = _stringify_metadata(
        {
            "conversation_id": conversation_id,
            **(dict(metadata) if metadata else {}),
        }
    )

    with propagate_attributes(
        session_id=conversation_id,
        metadata=propagated_metadata,
        tags=tags,
    ):
        yield