"""
Nora Observability Client
자동으로 AI 라이브러리 호출을 trace하고 API로 전송합니다.
"""
import time
import threading
from typing import Optional, Dict, Any, List
from datetime import datetime
import uuid
from contextlib import contextmanager
from contextvars import ContextVar

try:
    import requests
except ImportError:
    requests = None


# Context variables for trace grouping
_current_trace_group: ContextVar[Optional['TraceGroup']] = ContextVar('_current_trace_group', default=None)


class TraceGroup:
    """
    여러 LLM 호출을 하나의 논리적 그룹으로 묶는 컨텍스트.
    
    사용법:
        with nora.trace_group(name="multi_agent_pipeline"):
            # 이 블록 안의 모든 LLM 호출이 그룹으로 묶임
            response1 = client.chat.completions.create(...)
            response2 = client.chat.completions.create(...)
    """
    
    def __init__(self, name: str, metadata: Optional[Dict[str, Any]] = None):
        self.group_id = str(uuid.uuid4())
        self.name = name
        self.metadata = metadata or {}
        self.start_time = None
        self.end_time = None
        self.traces = []
        self._prev_auto_flush = None  # 이전 auto flush 상태 저장
        self._prev_trace_group = None  # 이전 trace_group 저장 (중첩 지원)
        
    def __enter__(self):
        self.start_time = time.time()
        # 이전 trace_group 저장 (중첩 지원)
        self._prev_trace_group = _current_trace_group.get()
        _current_trace_group.set(self)
        
        # 자동 플러시 비활성화 (trace_group 내부에서는 모아두기 위해)
        client = get_client()
        if client:
            self._prev_auto_flush = getattr(client, '_auto_flush_enabled', True)
            client._auto_flush_enabled = False
        
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.end_time = time.time()
        # 이전 trace_group 복원 (중첩 지원)
        _current_trace_group.set(self._prev_trace_group)
        
        # 자동 플러시 재개
        client = get_client()
        if client:
            if self._prev_auto_flush is not None:
                client._auto_flush_enabled = self._prev_auto_flush
        
        # 그룹 요약 정보 생성
        if self.traces:
            if client:
                group_summary = {
                    "group_id": self.group_id,
                    "group_name": self.name,
                    "trace_count": len(self.traces),
                    "total_duration": self.end_time - self.start_time,
                    "metadata": self.metadata,
                }
                # 각 trace에 그룹 정보 추가
                for trace in self.traces:
                    if trace.get("metadata") is None:
                        trace["metadata"] = {}
                    trace["metadata"]["trace_group"] = {
                        "id": self.group_id,
                        "name": self.name,
                    }
        
        return False  # 예외를 재발생시킴


class NoraClient:
    """
    Nora Observability 클라이언트
    
    Trace 데이터를 수집하고 배치로 API에 전송합니다.
    """
    
    def __init__(
        self,
        api_key: str,
        api_url: str = "https://api.nora.ai/v1/traces",
        batch_size: int = 10,
        flush_interval: float = 5.0,
    ):
        """
        Args:
            api_key: Nora API 키
            api_url: Trace 데이터를 전송할 API 엔드포인트 URL
            batch_size: 한 번에 전송할 trace 개수 (기본값: 10)
            flush_interval: 자동 전송 간격(초) (기본값: 5.0)
        """
        self.api_key = api_key
        self.api_url = api_url
        self.enabled = True
        self._auto_flush_enabled = True  # trace_group에서 제어 가능
        
        self._traces: List[Dict[str, Any]] = []
        self._lock = threading.Lock()
        self._batch_size = batch_size
        self._flush_interval = flush_interval
        self._last_flush = time.time()
        
    def trace(
        self,
        provider: str,
        model: str,
        prompt: Optional[str] = None,
        response: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
        start_time: Optional[float] = None,
        end_time: Optional[float] = None,
        tokens_used: Optional[int] = None,
        error: Optional[str] = None,
        finish_reason: Optional[str] = None,
        response_id: Optional[str] = None,
        system_fingerprint: Optional[str] = None,
        tool_calls: Optional[List[Dict[str, Any]]] = None,
        **extra_fields,
    ) -> None:
        """
        Trace 데이터를 수집합니다.
        
        Args:
            provider: AI 제공자 (openai, anthropic, etc.)
            model: 사용된 모델 이름
            prompt: 입력 프롬프트
            response: 응답 내용
            metadata: 추가 메타데이터
            start_time: 요청 시작 시간 (timestamp)
            end_time: 요청 종료 시간 (timestamp)
            tokens_used: 사용된 토큰 수
            error: 에러 메시지 (있는 경우)
            finish_reason: 완료 이유 (stop, length, tool_calls, etc.)
            response_id: API 응답 ID
            system_fingerprint: 시스템 fingerprint
            tool_calls: Tool/Function calls 정보
            **extra_fields: 추가 필드 (확장성)
        """
        if not self.enabled:
            return
        
        trace_data = {
            "id": str(uuid.uuid4()),
            "timestamp": datetime.utcnow().isoformat(),
            "provider": provider,
            "model": model,
            "prompt": prompt,
            "response": response,
            "metadata": metadata or {},
            "start_time": start_time,
            "end_time": end_time,
            "duration": (end_time - start_time) if (start_time and end_time) else None,
            "tokens_used": tokens_used,
            "error": error,
            "finish_reason": finish_reason,
            "response_id": response_id,
            "system_fingerprint": system_fingerprint,
            "tool_calls": tool_calls,
        }
        
        # 추가 필드 병합
        trace_data.update(extra_fields)
        
        # 현재 활성화된 trace group 정보 추가
        current_group = get_current_trace_group()
        
        if current_group:
            if trace_data["metadata"] is None:
                trace_data["metadata"] = {}
            trace_data["metadata"]["trace_group"] = {
                "id": current_group.group_id,
                "name": current_group.name,
            }
            current_group.traces.append(trace_data)
        
        with self._lock:
            self._traces.append(trace_data)
            
            # trace_group 내부에서는 자동 플러시 비활성화
            if not self._auto_flush_enabled:
                return
            
            # 배치 크기나 시간 간격에 따라 자동 전송
            should_flush = (
                len(self._traces) >= self._batch_size or
                (time.time() - self._last_flush) >= self._flush_interval
            )
            
            if should_flush:
                self._flush()
    
    def _flush(self) -> None:
        """수집된 trace 데이터를 API로 전송합니다."""
        if not self._traces:
            return
        
        if not requests:
            # requests가 없으면 경고 출력 (한 번만)
            if not hasattr(self, '_warned_no_requests'):
                print("[Nora] Warning: 'requests' library not found. Install it to send traces.")
                self._warned_no_requests = True
            return
        
        traces_to_send = self._traces.copy()
        self._traces.clear()
        self._last_flush = time.time()
        
        # 비동기로 전송 (메인 스레드 블로킹 방지)
        thread = threading.Thread(
            target=self._send_traces,
            args=(traces_to_send,),
            daemon=True
        )
        thread.start()
    
    def _send_traces(self, traces: List[Dict[str, Any]]) -> None:
        """실제 API로 trace 데이터를 전송합니다."""
        try:
            headers = {
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json",
            }
            
            payload = {"traces": traces}
            
            response = requests.post(
                self.api_url,
                json=payload,
                headers=headers,
                timeout=10
            )
            
            if response.status_code not in (200, 201):
                print(f"[Nora] Warning: Failed to send traces (status: {response.status_code})")
                
        except requests.exceptions.RequestException as e:
            # 네트워크 에러는 조용히 처리 (사용자 코드에 영향 없음)
            print(f"[Nora] Error sending traces: {str(e)}")
        except Exception as e:
            # 기타 예상치 못한 에러
            print(f"[Nora] Unexpected error: {str(e)}")
    
    def flush(self) -> None:
        """수동으로 trace 데이터를 즉시 전송합니다."""
        with self._lock:
            self._flush()
    
    def disable(self) -> None:
        """Trace 기능을 비활성화합니다."""
        self.flush()  # 비활성화 전에 남은 데이터 전송
        self.enabled = False
    
    def enable(self) -> None:
        """Trace 기능을 활성화합니다."""
        self.enabled = True
    
    def trace_group(self, name: str, metadata: Optional[Dict[str, Any]] = None) -> TraceGroup:
        """
        여러 LLM 호출을 하나의 논리적 그룹으로 묶습니다.
        
        Args:
            name: 그룹 이름
            metadata: 그룹 메타데이터
            
        Returns:
            TraceGroup context manager
            
        예제:
            >>> with client.trace_group("multi_agent_workflow"):
            ...     response1 = openai_client.chat.completions.create(...)
            ...     response2 = openai_client.chat.completions.create(...)
        """
        return TraceGroup(name=name, metadata=metadata)
    
    def find_traces_by_group(self, group_name: str) -> List[Dict[str, Any]]:
        """
        특정 trace group 이름으로 수집된 모든 traces를 검색합니다.
        
        Args:
            group_name: 검색할 trace group 이름
            
        Returns:
            매칭되는 trace들의 리스트
            
        예제:
            >>> traces = client.find_traces_by_group("multi_agent_pipeline")
            >>> for trace in traces:
            ...     print(f"Model: {trace['model']}, Tokens: {trace['tokens_used']}")
        """
        matching_traces = []
        with self._lock:
            for trace in self._traces:
                group_info = trace.get("metadata", {}).get("trace_group", {})
                if group_info.get("name") == group_name:
                    matching_traces.append(trace)
        return matching_traces
    
    def find_traces_by_group_id(self, group_id: str) -> List[Dict[str, Any]]:
        """
        특정 trace group ID로 수집된 모든 traces를 검색합니다.
        
        Args:
            group_id: 검색할 trace group ID
            
        Returns:
            매칭되는 trace들의 리스트
        """
        matching_traces = []
        with self._lock:
            for trace in self._traces:
                group_info = trace.get("metadata", {}).get("trace_group", {})
                if group_info.get("id") == group_id:
                    matching_traces.append(trace)
        return matching_traces
    
    def get_trace_groups(self) -> List[Dict[str, Any]]:
        """
        현재 수집된 모든 trace group 정보를 반환합니다.
        
        Returns:
            Unique한 trace group 정보 리스트 (id, name, trace_count, total_tokens, total_duration, trace_ids, metadata)
            
        예제:
            >>> groups = client.get_trace_groups()
            >>> for group in groups:
            ...     print(f"Group: {group['name']}, Traces: {group['trace_count']}")
        """
        groups_dict = {}
        with self._lock:
            for trace in self._traces:
                group_info = trace.get("metadata", {}).get("trace_group", {})
                if group_info:
                    group_id = group_info.get("id")
                    if group_id and group_id not in groups_dict:
                        groups_dict[group_id] = {
                            "id": group_id,
                            "name": group_info.get("name"),
                            "metadata": group_info.get("metadata", {}),
                            "trace_count": 0,
                            "total_tokens": 0,
                            "total_duration": 0.0,
                            "trace_ids": [],
                        }
                    if group_id:
                        groups_dict[group_id]["trace_count"] += 1
                        groups_dict[group_id]["total_tokens"] += trace.get("tokens_used", 0)
                        groups_dict[group_id]["trace_ids"].append(trace.get("id"))
                        duration = trace.get("duration")
                        if duration:
                            groups_dict[group_id]["total_duration"] += duration
        
        return list(groups_dict.values())



# 전역 클라이언트 인스턴스
_client: Optional[NoraClient] = None


def get_client() -> Optional[NoraClient]:
    """전역 클라이언트 인스턴스를 반환합니다."""
    return _client


def set_client(client: NoraClient) -> None:
    """전역 클라이언트 인스턴스를 설정합니다."""
    global _client
    _client = client


def get_current_trace_group() -> Optional[TraceGroup]:
    """현재 활성화된 trace group을 반환합니다."""
    return _current_trace_group.get()

