import os
import logging
from typing import Any, List, Optional

from langfuse.langchain import CallbackHandler as BaseCallbackHandler
from langchain_core.callbacks import BaseCallbackHandler as LangChainCallbackHandler
from langchain_core.messages import AIMessage
from langchain_core.outputs import LLMResult

logger = logging.getLogger(__name__)


def _is_enabled() -> bool:
    return os.getenv("LANGFUSE_ENABLED", "true").lower() == "true"


class CaptureAIMessageCallback(LangChainCallbackHandler):
    """LLM 호출 시 원본 AIMessage를 캡처해 usage(캐시 포함) 확인용으로 사용."""

    def __init__(self, storage: List[Any]):
        self.storage = storage

    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
        for gen_list in (response.generations or []):
            for gen in gen_list:
                if isinstance(gen.message, AIMessage):
                    self.storage.append(gen.message)
                    return


def get_langfuse_callback(
    trace_id: Optional[str] = None,
    session_id: Optional[str] = None,
    user_id: Optional[str] = None,
    parent_observation_id: Optional[str] = None,
    trace_name: Optional[str] = None,
):
    """
    Langfuse CallbackHandler 생성. 환경 변수 LANGFUSE_* 사용.

    - Core Agent: trace_id + trace_name="사용자 요청" 사용.
    - Sub Agent: trace_id + parent_observation_id 사용 (config.metadata에서 전달).
    """
    if not _is_enabled():
        return None
    try:
        public_key = os.getenv("LANGFUSE_PUBLIC_KEY")
        secret_key = os.getenv("LANGFUSE_SECRET_KEY")
        host = os.getenv("LANGFUSE_HOST", "https://langfuse.bai.biztechi.biz")
        if not public_key or not secret_key:
            logger.debug("LANGFUSE_PUBLIC_KEY or LANGFUSE_SECRET_KEY not set, skipping callback")
            return None
        os.environ["LANGFUSE_PUBLIC_KEY"] = public_key
        os.environ["LANGFUSE_SECRET_KEY"] = secret_key
        os.environ["LANGFUSE_HOST"] = host

        if trace_id:
            trace_ctx = {"trace_id": trace_id}
            if trace_name:
                trace_ctx["name"] = trace_name
            if parent_observation_id:
                trace_ctx["parent_span_id"] = parent_observation_id
            handler = BaseCallbackHandler(trace_context=trace_ctx)
        else:
            handler = BaseCallbackHandler()

        return handler
    except Exception as e:
        logger.error("Callback 생성 실패: %s", e)
        return None
