from sawsi.aws import shared
from sawsi.aws.s3_vector import wrapper
from typing import Optional, Literal
import json


class S3VectorAPI:
    """
    Amazon S3 Vectors 를 활용하는 벡터 저장/검색 API
    """
    def __init__(self, vector_bucket_name, credentials=None, region=shared.DEFAULT_REGION):
        """
        :param vector_bucket_name: 벡터 버킷 이름
        :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.vector_bucket_name = vector_bucket_name
        self.s3_vector = wrapper.S3Vector(self.boto3_session, region=region)

    def init_vector_bucket(self):
        """
        벡터 버킷 생성
        """
        try:
            return self.s3_vector.create_vector_bucket(self.vector_bucket_name)
        except Exception as ex:
            print(ex)

    def delete_vector_bucket(self):
        return self.s3_vector.delete_vector_bucket(self.vector_bucket_name)

    def get_vector_bucket(self):
        return self.s3_vector.get_vector_bucket(self.vector_bucket_name)

    def create_index(self, index_name, dimension, distance_metric='cosine', metadata_schema=None):
        """
        벡터 인덱스 생성
        :param index_name: 인덱스 이름
        :param dimension: 벡터 차원 수
        :param distance_metric: 거리 메트릭 ('cosine', 'euclidean', 'dotproduct')
        :param metadata_schema: 메타데이터 스키마 (dict)
        """
        return self.s3_vector.create_vector_index(
            self.vector_bucket_name, index_name, dimension, distance_metric, metadata_schema
        )

    def delete_index(self, index_name):
        return self.s3_vector.delete_vector_index(self.vector_bucket_name, index_name)

    def list_indexes(self, max_results=100, next_token=None):
        return self.s3_vector.list_vector_indexes(self.vector_bucket_name, max_results, next_token)

    def get_index(self, index_name):
        return self.s3_vector.get_vector_index(self.vector_bucket_name, index_name)

    def put_vectors(self, index_name, vectors):
        """
        벡터 저장
        :param index_name: 인덱스 이름
        :param vectors: [{'key': 'id1', 'data': {'float32': [0.1, 0.2, ...]}, 'metadata': {...}}, ...]
        """
        return self.s3_vector.put_vectors(self.vector_bucket_name, index_name, vectors)

    def put_vector(self, index_name, key, vector, metadata=None):
        """
        단일 벡터 저장
        :param index_name: 인덱스 이름
        :param key: 벡터 ID
        :param vector: 벡터 데이터 (list[float])
        :param metadata: 메타데이터 (dict)
        """
        item = {
            'key': key,
            'data': {'float32': vector},
        }
        if metadata:
            item['metadata'] = metadata
        return self.s3_vector.put_vectors(self.vector_bucket_name, index_name, [item])

    def get_vectors(self, index_name, keys):
        """
        벡터 조회
        :param index_name: 인덱스 이름
        :param keys: 벡터 ID 리스트
        """
        return self.s3_vector.get_vectors(self.vector_bucket_name, index_name, keys)

    def delete_vectors(self, index_name, keys):
        return self.s3_vector.delete_vectors(self.vector_bucket_name, index_name, keys)

    def query(self, index_name, query_vector, top_k=10, filter_expression=None):
        """
        벡터 유사도 검색
        :param index_name: 인덱스 이름
        :param query_vector: 쿼리 벡터 (list[float])
        :param top_k: 반환할 결과 수
        :param filter_expression: 메타데이터 필터
        :return: 유사한 벡터 목록
        """
        return self.s3_vector.query_vectors(
            self.vector_bucket_name, index_name, query_vector, top_k, filter_expression
        )

    def put_vectors_batch(self, index_name, keys, vectors, metadata_list=None, batch_size=100):
        """
        벡터 배치 저장
        :param index_name: 인덱스 이름
        :param keys: 벡터 ID 리스트
        :param vectors: 벡터 데이터 리스트 (list[list[float]])
        :param metadata_list: 메타데이터 리스트 (list[dict])
        :param batch_size: 배치 크기
        """
        items = []
        for i, (key, vector) in enumerate(zip(keys, vectors)):
            item = {
                'key': key,
                'data': {'float32': vector},
            }
            if metadata_list and i < len(metadata_list):
                item['metadata'] = metadata_list[i]
            items.append(item)

        results = []
        for i in range(0, len(items), batch_size):
            batch = items[i:i + batch_size]
            result = self.s3_vector.put_vectors(self.vector_bucket_name, index_name, batch)
            results.append(result)
        return results


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