from sawsi.aws import shared
from sawsi.aws.bedrock.api import BedrockAPI
from sawsi.aws.s3_vector.api import S3VectorAPI
from typing import Optional
from secrets import token_urlsafe


DEFAULT_EMBED_MODEL = 'amazon.titan-embed-text-v2:0'
DEFAULT_DIMENSION = 1024


class VectorKBAPI:
    """
    텍스트 기반 벡터 지식베이스 API
    Bedrock 임베딩 + S3 Vectors 를 조합하여
    텍스트만 넣으면 자동으로 임베딩 → 저장 → 검색까지 처리합니다.

    사용 예:
        kb = VectorKBAPI('my-kb-bucket', 'my-index')
        kb.init()

        # 텍스트 저장
        kb.add('doc-1', '서울은 대한민국의 수도입니다.', metadata={'source': 'wiki'})

        # 검색
        results = kb.search('한국의 수도는?', top_k=3)

        # 배치 저장
        kb.add_batch([
            {'key': 'doc-2', 'text': '부산은 항구도시입니다.'},
            {'key': 'doc-3', 'text': '제주도는 섬입니다.'},
        ])
    """
    def __init__(
        self,
        vector_bucket_name,
        index_name,
        embed_model_id=DEFAULT_EMBED_MODEL,
        dimension=DEFAULT_DIMENSION,
        distance_metric='cosine',
        credentials=None,
        region=shared.DEFAULT_REGION,
    ):
        """
        :param vector_bucket_name: S3 Vectors 버킷 이름
        :param index_name: 벡터 인덱스 이름
        :param embed_model_id: Bedrock 임베딩 모델 ID (기본: Titan Embed Text v2)
        :param dimension: 벡터 차원 수 (기본: 1024)
        :param distance_metric: 거리 메트릭 ('cosine', 'euclidean', 'dotproduct')
        :param credentials: AWS 자격 증명
        :param region: AWS 리전
        """
        self.index_name = index_name
        self.embed_model_id = embed_model_id
        self.dimension = dimension
        self.distance_metric = distance_metric
        self.bedrock = BedrockAPI(credentials=credentials, region=region)
        self.vector_store = S3VectorAPI(vector_bucket_name, credentials=credentials, region=region)

    def init(self, metadata_schema=None):
        """
        벡터 버킷 및 인덱스를 초기화합니다. (최초 1회)
        :param metadata_schema: 메타데이터 스키마 (dict), 예: {"type": "object", "properties": {"source": {"type": "string"}}}
        """
        self.vector_store.init_vector_bucket()
        try:
            self.vector_store.create_index(
                self.index_name, self.dimension, self.distance_metric, metadata_schema
            )
        except Exception as ex:
            print(ex)

    def _embed(self, text):
        """텍스트를 벡터로 변환"""
        return self.bedrock.get_embeddings(
            self.embed_model_id, text, dimensions=self.dimension
        )

    def add(self, key, text, metadata=None):
        """
        텍스트를 임베딩하여 벡터 저장소에 저장합니다.
        :param key: 문서 ID
        :param text: 저장할 텍스트
        :param metadata: 메타데이터 (dict)
        :return: 저장 결과
        """
        vector = self._embed(text)
        meta = metadata or {}
        meta['text'] = text
        return self.vector_store.put_vector(self.index_name, key, vector, metadata=meta)

    def add_auto_key(self, text, metadata=None):
        """
        키를 자동 생성하여 텍스트를 저장합니다.
        :param text: 저장할 텍스트
        :param metadata: 메타데이터 (dict)
        :return: (생성된 key, 저장 결과)
        """
        key = token_urlsafe(16)
        result = self.add(key, text, metadata)
        return key, result

    def add_batch(self, items, batch_size=10):
        """
        여러 텍스트를 배치로 임베딩하여 저장합니다.
        :param items: [{'key': 'doc-1', 'text': '...', 'metadata': {...}}, ...]
                      metadata 는 선택사항
        :param batch_size: 한 번에 저장할 벡터 수
        :return: 저장 결과 리스트
        """
        keys = []
        vectors = []
        metadata_list = []

        for item in items:
            key = item.get('key', token_urlsafe(16))
            text = item['text']
            meta = item.get('metadata', {})
            meta['text'] = text

            keys.append(key)
            vectors.append(self._embed(text))
            metadata_list.append(meta)

        return self.vector_store.put_vectors_batch(
            self.index_name, keys, vectors, metadata_list, batch_size
        )

    def search(self, query_text, top_k=5, filter_expression=None):
        """
        텍스트로 유사도 검색을 수행합니다.
        :param query_text: 검색 쿼리 텍스트
        :param top_k: 반환할 결과 수
        :param filter_expression: 메타데이터 필터
        :return: 유사한 벡터 목록
        """
        query_vector = self._embed(query_text)
        return self.vector_store.query(
            self.index_name, query_vector, top_k, filter_expression
        )

    def get(self, keys):
        """
        키로 벡터를 조회합니다.
        :param keys: 벡터 ID 리스트
        """
        if isinstance(keys, str):
            keys = [keys]
        return self.vector_store.get_vectors(self.index_name, keys)

    def delete(self, keys):
        """
        키로 벡터를 삭제합니다.
        :param keys: 벡터 ID 리스트
        """
        if isinstance(keys, str):
            keys = [keys]
        return self.vector_store.delete_vectors(self.index_name, keys)


if __name__ == '__main__':
    kb = VectorKBAPI('my-kb-bucket', 'my-index')
    kb.add('1', 'SOME TEXT')
    kb.add('2', 'HELLO')
    r = kb.search('SOME')
    print(r)
