from sawsi.aws import shared
from sawsi.aws.bedrock import wrapper
from typing import Optional


class BedrockAPI:
    """
    AWS Bedrock 런타임 API
    """
    def __init__(self, credentials=None, region=shared.DEFAULT_REGION):
        """
        :param credentials: {
            "aws_access_key_id": "str",
            "aws_secret_access_key": "str",
            "region_name": "str",
            "profile_name": "str",
        }
        """
        self.boto3_session = shared.get_boto_session(credentials)
        self.bedrock = wrapper.Bedrock(self.boto3_session, region=region)

    def invoke_model(self, model_id, body, content_type='application/json', accept='application/json'):
        """
        모델 직접 호출 (로우레벨)
        """
        return self.bedrock.invoke_model(model_id, body, content_type, accept)

    def converse(self, model_id, messages, system=None, max_tokens=None, temperature=None, top_p=None, tool_config=None):
        """
        Converse API 를 통한 모델 호출
        :param model_id: 모델 ID (예: 'anthropic.claude-sonnet-4-20250514')
        :param messages: [{'role': 'user', 'content': [{'text': '...'}]}]
        :param system: 시스템 프롬프트 (str 또는 [{'text': '...'}])
        :param max_tokens: 최대 토큰 수
        :param temperature: 온도
        :param top_p: top_p
        :param tool_config: 도구 설정
        """
        system_param = None
        if system:
            if isinstance(system, str):
                system_param = [{'text': system}]
            else:
                system_param = system

        inference_config = {}
        if max_tokens:
            inference_config['maxTokens'] = max_tokens
        if temperature is not None:
            inference_config['temperature'] = temperature
        if top_p is not None:
            inference_config['topP'] = top_p

        return self.bedrock.converse(
            model_id, messages,
            system=system_param,
            inference_config=inference_config if inference_config else None,
            tool_config=tool_config,
        )

    def converse_stream(self, model_id, messages, system=None, max_tokens=None, temperature=None, top_p=None, tool_config=None):
        """
        Converse API 스트리밍 호출
        """
        system_param = None
        if system:
            if isinstance(system, str):
                system_param = [{'text': system}]
            else:
                system_param = system

        inference_config = {}
        if max_tokens:
            inference_config['maxTokens'] = max_tokens
        if temperature is not None:
            inference_config['temperature'] = temperature
        if top_p is not None:
            inference_config['topP'] = top_p

        return self.bedrock.converse_stream(
            model_id, messages,
            system=system_param,
            inference_config=inference_config if inference_config else None,
            tool_config=tool_config,
        )

    def chat(self, model_id, user_message, system=None, max_tokens=1024, temperature=None):
        """
        간단한 채팅 메시지 전송
        :param model_id: 모델 ID
        :param user_message: 사용자 메시지 (str)
        :param system: 시스템 프롬프트 (str)
        :param max_tokens: 최대 토큰 수
        :param temperature: 온도
        :return: 응답 텍스트 (str)
        """
        messages = [
            {'role': 'user', 'content': [{'text': user_message}]}
        ]
        response = self.converse(model_id, messages, system=system, max_tokens=max_tokens, temperature=temperature)
        return response['output']['message']['content'][0]['text']

    def get_embeddings(self, model_id, input_text, dimensions=None, normalize=True):
        """
        임베딩 벡터 생성 (Amazon Titan Embeddings 등)
        :param model_id: 임베딩 모델 ID (예: 'amazon.titan-embed-text-v2:0')
        :param input_text: 입력 텍스트
        :param dimensions: 벡터 차원 수
        :param normalize: 정규화 여부
        :return: 임베딩 벡터 (list[float])
        """
        body = {
            'inputText': input_text,
        }
        if dimensions:
            body['dimensions'] = dimensions
        if normalize is not None:
            body['normalize'] = normalize
        response = self.invoke_model(model_id, body)
        return response.get('embedding', response)

    def generate_stream_text(self, model_id, messages, system=None, max_tokens=1024, temperature=None):
        """
        스트리밍 텍스트 생성 제너레이터
        :return: 텍스트 청크를 yield
        """
        response = self.converse_stream(model_id, messages, system=system, max_tokens=max_tokens, temperature=temperature)
        stream = response.get('stream', [])
        for event in stream:
            if 'contentBlockDelta' in event:
                delta = event['contentBlockDelta'].get('delta', {})
                if 'text' in delta:
                    yield delta['text']


if __name__ == '__main__':
    api = BedrockAPI()
