"""
Bedrock API 호출 수집기
boto3를 통한 Bedrock API 호출을 추적함

역할: boto3 bedrock-runtime 클라이언트의 API 호출을 자동으로 추적함
사용: MonitoringSDK.init()에서 자동 등록됨
"""

import time
import json
import io
from datetime import datetime
from typing import Any, Dict, Optional
from ..types.interfaces import SpanEntry
from ..transport.batch_sender import BatchSender
from ..tracing.trace_context import get_trace_context
from ..tracing.id_generator import generate_span_id


class BedrockCollector:
    """
    Bedrock API 호출을 추적하는 수집기
    """

    def __init__(self, batch_sender: BatchSender, service_name: str, environment: str):
        """
        BedrockCollector 초기화

        Args:
            batch_sender: 배치 전송자
            service_name: 서비스 이름
            environment: 환경 설정
        """
        self.batch_sender = batch_sender
        self.service_name = service_name
        self.environment = environment

    def setup_bedrock_instrumentation(self) -> None:
        """
        boto3 bedrock-runtime 클라이언트에 instrumentation을 설정함

        역할: bedrock-runtime API 호출 시 자동으로 추적하도록 설정
        """
        try:
            import boto3
            from botocore.client import BaseClient

            # boto3 client의 _make_api_call 메서드를 패치
            original_make_api_call = BaseClient._make_api_call

            def patched_make_api_call(client_self, operation_name, api_params):
                # bedrock-runtime 서비스만 추적
                service_name = client_self._service_model.service_name

                if service_name != "bedrock-runtime":
                    return original_make_api_call(
                        client_self, operation_name, api_params
                    )

                trace_context = get_trace_context()
                if not trace_context:
                    return original_make_api_call(
                        client_self, operation_name, api_params
                    )

                start_time = time.time()
                status = "ERROR"
                response = None
                error = None

                try:
                    response = original_make_api_call(
                        client_self, operation_name, api_params
                    )
                    status = "OK"
                    return response
                except Exception as e:
                    status = "ERROR"
                    error = e
                    raise
                finally:
                    duration = (time.time() - start_time) * 1000
                    self._record_bedrock_span(
                        operation_name,
                        api_params,
                        response,
                        duration,
                        status,
                        trace_context.trace_id,
                        trace_context.span_id,
                        error,
                    )

            BaseClient._make_api_call = patched_make_api_call

        except ImportError:
            # boto3가 설치되지 않은 경우 무시
            pass

    def _record_bedrock_span(
        self,
        operation_name: str,
        api_params: Dict[str, Any],
        response: Optional[Dict[str, Any]],
        duration: float,
        status: str,
        trace_id: str,
        parent_span_id: str,
        error: Optional[Exception],
    ) -> None:
        """
        Bedrock API 호출 span을 기록함

        Args:
            operation_name: Bedrock API 작업 이름 (예: InvokeModel, InvokeModelWithResponseStream)
            api_params: API 파라미터
            response: API 응답
            duration: 소요 시간 (밀리초)
            status: 상태 (OK 또는 ERROR)
            trace_id: Trace ID
            parent_span_id: Parent Span ID
            error: 에러 (있는 경우)
        """
        span_id = generate_span_id()
        timestamp = datetime.utcnow().isoformat() + "Z"

        # 모델 ID 추출
        model_id = api_params.get("modelId", "unknown")

        # 토큰 정보 추출 (응답이 있고 status가 OK인 경우)
        input_tokens = None
        output_tokens = None

        if response and status == "OK":
            # Bedrock 응답에서 토큰 정보 추출
            try:
                # InvokeModel 응답의 경우
                if "body" in response:
                    # body를 읽기 전에 원본을 보존하기 위해 내용을 복사
                    original_body = response["body"]
                    body_bytes = original_body.read()

                    # 읽은 내용으로 JSON 파싱하여 토큰 정보 추출
                    # body_bytes는 bytes이므로 decode 필요
                    if isinstance(body_bytes, bytes):
                        body_str = body_bytes.decode("utf-8")
                    else:
                        body_str = body_bytes
                    response_body = json.loads(body_str)

                    # Claude 모델의 경우
                    if "usage" in response_body:
                        input_tokens = response_body["usage"].get("input_tokens")
                        output_tokens = response_body["usage"].get("output_tokens")
                    # 다른 모델의 경우 여기에 추가

                    # body를 읽은 후 원본을 복원하여 bedrock_service에서 사용할 수 있도록 함
                    # io.BytesIO를 사용하여 새로운 스트림 생성
                    from botocore.response import StreamingBody

                    response["body"] = StreamingBody(
                        io.BytesIO(body_bytes),
                        (
                            original_body._content_length
                            if hasattr(original_body, "_content_length")
                            else len(body_bytes)
                        ),
                    )
            except Exception:
                # 토큰 정보 추출 실패 시 무시
                pass

        span: SpanEntry = {
            "type": "span",
            "timestamp": timestamp,
            "service_name": self.service_name,
            "environment": self.environment,
            "trace_id": trace_id,
            "span_id": span_id,
            "parent_span_id": parent_span_id,
            "name": f"Bedrock {operation_name}",
            "kind": "CLIENT",
            "duration_ms": duration,
            "status": status,  # type: ignore
            "bedrock_model_id": model_id,
            "bedrock_operation": operation_name,
            "bedrock_input_tokens": input_tokens,
            "bedrock_output_tokens": output_tokens,
        }

        self.batch_sender.add(span)
