import os
import logging
from typing import Any, Dict, Optional, Tuple, Callable

import requests
import urllib3
from langfuse import get_client

logger = logging.getLogger(__name__)


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


def _get_env(key: str, default: str = "") -> str:
    return os.getenv(key, default)


# 개발 환경에서 Langfuse self-signed SSL 비활성화 (선택)
_env = _get_env("ENVIRONMENT", "development")
if _env == "development":
    urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
    _original_request = requests.Session.request
    _langfuse_host = _get_env("LANGFUSE_HOST", "").replace("https://", "").replace("http://", "").rstrip("/")

    def _patched_request(self, method, url, *args, **kwargs):
        if _langfuse_host and _langfuse_host in url:
            kwargs["verify"] = False
        if "langfuse.bai.biztechi.biz" in url:
            kwargs["verify"] = False
        return _original_request(self, method, url, *args, **kwargs)

    requests.Session.request = _patched_request


def _get_langfuse_client():
    if not _is_enabled():
        return None
    return get_client()


def create_child_span(
    trace_id: str, name: str, parent_observation_id: Optional[str] = None
):
    """자식 span 생성. parent_observation_id가 있으면 그 아래에, 없으면 trace 루트 아래에."""
    client = _get_langfuse_client()
    if not client:
        return None
    try:
        trace_ctx = {
            "trace_id": trace_id,
            **({"parent_span_id": parent_observation_id} if parent_observation_id else {}),
        }
        return client.start_as_current_observation(
            name=name, as_type="span", trace_context=trace_ctx
        )
    except Exception as e:
        logger.warning("Child span 생성 실패 (%s): %s", name, e)
        return None


def get_current_observation_id() -> Optional[str]:
    client = _get_langfuse_client()
    if not client:
        return None
    try:
        return client.get_current_observation_id()
    except Exception:
        return None


def get_usage_from_response(response: Any) -> Optional[Dict[str, Any]]:
    """AIMessage 등 LLM 응답에서 usage dict 추출 (input, output, cache_read_input_tokens, cache_write_input_tokens)."""
    usage_info = None
    if hasattr(response, "response_metadata") and response.response_metadata:
        usage_info = response.response_metadata.get("usage") or {}
    if not usage_info and getattr(response, "usage_metadata", None):
        usage_info = getattr(response, "usage_metadata") or {}
    if not usage_info:
        return None
    cache_read = (
        usage_info.get("cache_read_input_tokens")
        or usage_info.get("cacheReadInputTokens")
        or 0
    )
    cache_write = (
        usage_info.get("cache_write_input_tokens")
        or usage_info.get("cacheWriteInputTokens")
        or 0
    )
    details = usage_info.get("input_token_details") or {}
    if not cache_read and not cache_write:
        cache_read = details.get("cache_read") or details.get("cacheRead") or 0
        cache_write = details.get("cache_creation") or details.get("cacheCreation") or 0
    return {
        "input": usage_info.get("input_tokens")
        or usage_info.get("inputTokens")
        or usage_info.get("prompt_tokens", 0),
        "output": usage_info.get("output_tokens")
        or usage_info.get("outputTokens")
        or usage_info.get("completion_tokens", 0),
        "cache_read_input_tokens": cache_read,
        "cache_write_input_tokens": cache_write,
    }


def _extract_llm_info(
    llm: Any,
    response: Any,
    output_extractor: Optional[Callable] = None,
) -> Tuple[str, Optional[Dict[str, int]], Dict[str, Any], str]:
    if output_extractor:
        output = output_extractor(response)
    elif hasattr(response, "content"):
        output = response.content
    else:
        output = str(response)
    usage = get_usage_from_response(response)
    model_parameters = {
        "temperature": getattr(llm, "temperature", None),
        "max_tokens": getattr(llm, "max_tokens", None),
    }
    model = getattr(llm, "model_id", "unknown")
    return output, usage, model_parameters, model


def log_llm_generation(
    state: Any,
    node_name: str,
    llm: Any,
    messages: list,
    response: Any,
    output_extractor: Optional[Callable] = None,
) -> None:
    """
    LLM generation 로깅. state는 .get("context", {}).get("trace_id")를 지원하는 객체(dict 등).
    """
    client = _get_langfuse_client()
    if not client:
        return

    context = state.get("context", {}) if hasattr(state, "get") else {}
    trace_id = context.get("trace_id") if isinstance(context, dict) else None
    if not trace_id:
        logger.warning("trace_id를 찾을 수 없어 Langfuse 기록을 건너뜁니다 (%s)", node_name)
        return

    output, usage, model_parameters, model = _extract_llm_info(
        llm, response, output_extractor
    )

    try:
        parent_observation_id = get_current_observation_id()
        trace_ctx = {
            "trace_id": trace_id,
            **({"parent_span_id": parent_observation_id} if parent_observation_id else {}),
        }
        usage_details = None
        if usage:
            usage_details = {
                "input": usage.get("input", 0),
                "output": usage.get("output", 0),
                "cache_read_input_tokens": usage.get("cache_read_input_tokens", 0),
                "cache_creation_input_tokens": usage.get("cache_write_input_tokens", 0),
            }

        agent_type = _get_env("AGENT_TYPE", "Core")
        with client.start_as_current_observation(
            name=node_name,
            as_type="generation",
            model=model,
            input=str(messages),
            model_parameters=model_parameters or {},
            trace_context=trace_ctx,
            metadata={"agent_type": agent_type},
        ) as generation:
            generation.update(
                output=output,
                **({"usage_details": usage_details} if usage_details else {}),
            )
    except Exception as e:
        logger.warning("LLM generation 기록 실패 (%s): %s", node_name, e)


def log_llm_generation_span(
    node_name: str,
    trace_id: str,
    model: str,
    input_prompt: Optional[str] = None,
    input_messages: Optional[list] = None,
    output: str = "",
    usage: Optional[Dict[str, Any]] = None,
    model_parameters: Optional[Dict] = None,
    parent_observation_id: Optional[str] = None,
) -> None:
    """
    Sub Agent용 LLM generation 로깅 (명시적 인자).
    trace_id, model, input/output, usage를 직접 넘겨 노드 이름으로 비용 집계.
    """
    client = _get_langfuse_client()
    if not client:
        return
    if not trace_id:
        logger.warning("trace_id를 찾을 수 없어 Langfuse 기록을 건너뜁니다 (%s)", node_name)
        return
    if not model:
        logger.warning("model을 찾을 수 없어 Langfuse 기록을 건너뜁니다 (%s)", node_name)
        return

    input_value = str(input_prompt) if input_prompt else (str(input_messages) if input_messages else "")

    try:
        if not parent_observation_id:
            parent_observation_id = get_current_observation_id()
        trace_ctx = {"trace_id": trace_id}
        if parent_observation_id:
            trace_ctx["parent_span_id"] = parent_observation_id

        usage_details = None
        if usage:
            usage_details = {
                "input": usage.get("input", 0),
                "output": usage.get("output", 0),
                "cache_read_input_tokens": usage.get("cache_read_input_tokens", 0),
                "cache_creation_input_tokens": usage.get("cache_write_input_tokens", 0),
            }
            if usage_details["cache_read_input_tokens"] == 0 and usage_details["cache_creation_input_tokens"] == 0:
                usage_details = {"input": usage.get("input", 0), "output": usage.get("output", 0)}

        agent_type = _get_env("AGENT_TYPE", "Core")
        with client.start_as_current_observation(
            name=node_name,
            as_type="generation",
            model=model,
            input=input_value,
            model_parameters=model_parameters or {},
            trace_context=trace_ctx,
            metadata={"agent_type": agent_type},
        ) as generation:
            generation.update(
                output=output,
                **({"usage_details": usage_details} if usage_details else {}),
            )
    except Exception as e:
        logger.warning("LLM generation 기록 실패 (%s): %s", node_name, e)


def flush_langfuse() -> None:
    """대기 중인 데이터를 Langfuse로 전송."""
    client = _get_langfuse_client()
    if client:
        client.flush()
