from sawsi.aws import shared
import json


class S3Vector:
    def __init__(self, boto3_session, region=shared.DEFAULT_REGION):
        self.client = boto3_session.client('s3vectors', region_name=region)
        self.region = region

    def create_vector_bucket(self, vector_bucket_name, encryption_type='aws:s3vectors'):
        """
        벡터 버킷 생성
        """
        params = {
            'vectorBucketName': vector_bucket_name,
            'encryptionConfiguration': {
                'sseType': encryption_type,
            },
        }
        response = self.client.create_vector_bucket(**params)
        return response

    def delete_vector_bucket(self, vector_bucket_name):
        response = self.client.delete_vector_bucket(vectorBucketName=vector_bucket_name)
        return response

    def list_vector_buckets(self, max_results=100, next_token=None):
        params = {'maxResults': max_results}
        if next_token:
            params['nextToken'] = next_token
        response = self.client.list_vector_buckets(**params)
        return response

    def get_vector_bucket(self, vector_bucket_name):
        response = self.client.get_vector_bucket(vectorBucketName=vector_bucket_name)
        return response

    def create_vector_index(self, vector_bucket_name, vector_index_name, dimension, distance_metric='cosine', metadata_schema=None):
        """
        벡터 인덱스 생성
        :param vector_bucket_name: 벡터 버킷 이름
        :param vector_index_name: 벡터 인덱스 이름
        :param dimension: 벡터 차원 수
        :param distance_metric: 거리 메트릭 ('cosine', 'euclidean', 'dotproduct')
        :param metadata_schema: 메타데이터 스키마 (dict)
        """
        params = {
            'vectorBucketName': vector_bucket_name,
            'vectorIndexName': vector_index_name,
            'dimension': dimension,
            'distanceMetric': distance_metric,
        }
        if metadata_schema:
            params['metadataConfiguration'] = {
                'metadataSchema': json.dumps(metadata_schema),
            }
        response = self.client.create_vector_index(**params)
        return response

    def delete_vector_index(self, vector_bucket_name, vector_index_name):
        response = self.client.delete_vector_index(
            vectorBucketName=vector_bucket_name,
            vectorIndexName=vector_index_name,
        )
        return response

    def list_vector_indexes(self, vector_bucket_name, max_results=100, next_token=None):
        params = {
            'vectorBucketName': vector_bucket_name,
            'maxResults': max_results,
        }
        if next_token:
            params['nextToken'] = next_token
        response = self.client.list_vector_indexes(**params)
        return response

    def get_vector_index(self, vector_bucket_name, vector_index_name):
        response = self.client.get_vector_index(
            vectorBucketName=vector_bucket_name,
            vectorIndexName=vector_index_name,
        )
        return response

    def put_vectors(self, vector_bucket_name, vector_index_name, vectors):
        """
        벡터 저장
        :param vectors: [{'key': 'id1', 'data': {'float32': [0.1, 0.2, ...]}, 'metadata': {...}}, ...]
        """
        response = self.client.put_vectors(
            vectorBucketName=vector_bucket_name,
            vectorIndexName=vector_index_name,
            vectors=vectors,
        )
        return response

    def get_vectors(self, vector_bucket_name, vector_index_name, keys):
        """
        벡터 조회
        :param keys: ['id1', 'id2', ...]
        """
        response = self.client.get_vectors(
            vectorBucketName=vector_bucket_name,
            vectorIndexName=vector_index_name,
            keys=keys,
        )
        return response

    def delete_vectors(self, vector_bucket_name, vector_index_name, keys):
        response = self.client.delete_vectors(
            vectorBucketName=vector_bucket_name,
            vectorIndexName=vector_index_name,
            keys=keys,
        )
        return response

    def query_vectors(self, vector_bucket_name, vector_index_name, query_vector, top_k=10, filter_expression=None):
        """
        벡터 유사도 검색
        :param query_vector: 쿼리 벡터 (list[float])
        :param top_k: 반환할 결과 수
        :param filter_expression: 메타데이터 필터
        """
        params = {
            'vectorBucketName': vector_bucket_name,
            'vectorIndexName': vector_index_name,
            'queryVector': {'float32': query_vector},
            'topK': top_k,
        }
        if filter_expression:
            params['filter'] = filter_expression
        response = self.client.query_vectors(**params)
        return response
