
"""
검색엔진 개발 예시
"""
import regex
from sawsi.aws.dynamodb.api import DynamoDBAPI
from typing import Optional
from collections import Counter
from decimal import Decimal


def levenshtein_distance(s: str, t: str) -> int:
    """
    두 문자열 s와 t 사이의 Levenshtein 거리를 계산합니다.
    동적 프로그래밍을 이용한 기본 O(len(s)*len(t)) 구현입니다.
    """
    if s == t:
        return 0
    if len(s) == 0:
        return len(t)
    if len(t) == 0:
        return len(s)

    previous_row = list(range(len(t) + 1))
    for i, sc in enumerate(s, start=1):
        current_row = [i]
        for j, tc in enumerate(t, start=1):
            insert_cost = previous_row[j] + 1
            delete_cost = current_row[j - 1] + 1
            substitute_cost = previous_row[j - 1] + (0 if sc == tc else 1)
            current_row.append(min(insert_cost, delete_cost, substitute_cost))
        previous_row = current_row
    return previous_row[-1]


def levenshtein_ratio(s: str, t: str) -> float:
    """
    Levenshtein 거리를 바탕으로 두 문자열의 유사도 비율을 계산합니다.
    유사도 = (max_len - distance) / max_len, max_len이 0이면 1.0 반환.
    """
    max_len = max(len(s), len(t))
    if max_len == 0:
        return 1.0
    return (max_len - levenshtein_distance(s, t)) / max_len


class SearchEngine:
    def __init__(self, dynamo_db_api: DynamoDBAPI, table_name: str):
        self.table_name = table_name
        self.dynamo_db_api = dynamo_db_api

    def init_table(self):
        response = self.dynamo_db_api.create_table(
            self.table_name,
            partition_key='i',
            partition_key_type='S',
            sort_key='p',
            sort_key_type='N'
        )
        print('response:', response)
        self._create_gsi()
        return response

    def _create_gsi(self):
        response = self.dynamo_db_api.create_global_index(
            self.table_name,
            partition_key='g',
            partition_key_type='S',
            sort_key='s',
            sort_key_type='N',
        )
        return response

    @classmethod
    def _make_n_gram(cls, text: str, n: int = 2) -> list[str]:
        """
        문자열을 받아 n-gram 리스트를 반환하는 함수
        :param text: 입력 문자열
        :param n: n-gram의 n 값 (기본값: 2)
        :return: n-gram 리스트
        """
        if n < 1:
            raise ValueError("n은 1 이상의 정수여야 합니다.")
        return [text[i:i + n] for i in range(len(text) - n + 1)]

    @classmethod
    def _pre_process(cls, value: str) -> str:
        """
        특수문자 및 이모지를 제거하고, 앞에 ^ 부착
        :param value: 입력 문자열
        :return: 전처리된 문자열
        """
        v = regex.sub(r'[^\p{L}\p{N}]', '', value)  # 문자(Letter) 및 숫자(Number)만 남김
        v = f'^{v}'  # 앞에 ^ 부착
        return v

    def put_data(self, item_id: str, text: str):
        """
        레코드에 특정 필드의 text 에 대해 삽입
        field_value에서 특수문자 다 삭제하고 앞에 ^를 붙여서 한글자 검색 가능하게 변경
        :param item_id:
        :param text:
        :return:
        """
        v = self._pre_process(text)
        v_size = len(v)
        last_p = 0
        for idx, gram in enumerate(self._make_n_gram(v, 2)):
            self.dynamo_db_api.put_item(self.table_name, {
                'g': gram,  # 토큰
                'i': item_id,  # 가리키는 Item ID
                's': v_size,  # 총 문자열 사이즈
                'p': idx,  # 이 토큰의 gram 위치
            })
            last_p = idx
        # 입력된 총 문자열 사이즈보다 큰 것들은 다 삭제
        item_deleted = self.dynamo_db_api.delete_item_by_key(
            self.table_name,
            {
                'i': item_id,
                'p': last_p + 1
            }
        )
        item_deleted = item_deleted.get('Attributes', {})
        # print('first-item_deleted:', item_deleted)
        if item_deleted:  # 있으면, v_size 를 구해서 그만큼 삭제
            old_size = item_deleted['s']
            old_size = int(old_size)
            for idx in range(last_p + 2, old_size):
                item_deleted = self.dynamo_db_api.delete_item_by_key(
                    self.table_name,
                    {
                        'i': item_id,
                        'p': idx,
                    }
                )
                item_deleted = item_deleted.get('Attributes', {})
                # print('item_deleted:', item_deleted)

    def delete_data(self, item_id: str):
        """
        특정 Item 에 대해 검색 레코드 삭제
        :param item_id:
        :return:
        """
        # 일단 첫 토큰 삭제
        item_deleted = self.dynamo_db_api.delete_item_by_key(
            self.table_name,
            {
                'i': item_id,
                'p': 0,
            }
        )
        item_deleted = item_deleted.get('Attributes', {})
        # print('first-item_deleted:', item_deleted)
        if item_deleted:  # 있으면, v_size 를 구해서 그만큼 삭제
            old_size = item_deleted['s']
            old_size = int(old_size)
            for idx in range(1, old_size):
                item_deleted = self.dynamo_db_api.delete_item_by_key(
                    self.table_name,
                    {
                        'i': item_id,
                        'p': idx,
                    }
                )
                item_deleted = item_deleted.get('Attributes', {})
                # print('item_deleted:', item_deleted)


    @staticmethod
    def _convert_decimal(obj):
        """
        객체 내의 Decimal 타입을 int/float 등 JSON 직렬화 가능한 타입으로 변환.
        """
        if isinstance(obj, list):
            return [SearchEngine._convert_decimal(i) for i in obj]
        elif isinstance(obj, dict):
            return {k: SearchEngine._convert_decimal(v) for k, v in obj.items()}
        elif isinstance(obj, Decimal):
            # 정수인지 여부에 따라 int 또는 float로 변환
            return int(obj) if obj % 1 == 0 else float(obj)
        else:
            return obj

    def search(self, keyword: str, start_key: Optional[dict], limit: int = 10) -> tuple[list[str], Optional[dict]]:
        """
        개선된 n-gram 기반 검색.

        - 검색어를 전처리 후 2-gram 리스트를 생성합니다.
        - 각 n-gram에 대해 DynamoDB에서 항목들을 조회하여, 후보 문서별로
          (검색어 내 토큰 인덱스, 문서 내 토큰 위치) 매칭 정보를 수집합니다.
        - 각 후보에 대해 다음 세 가지 점수를 산출합니다.
            1. Coverage: 검색어 토큰 중 몇 개가 매칭되었는지 (고유 토큰 비율)
            2. Continuity: 검색어 토큰들이 연속해서 매칭된 최대 run 길이의 비율
            3. Position Score: 1 / (1 + 평균 매칭 위치)

        최종 점수 = (α * coverage) + (β * continuity) + (γ * position_score)

        문서 길이에 의한 영향을 배제하여, 단순히 검색어와의 매칭 패턴에만 집중합니다.

        정렬은 최종 점수 내림차순, tie-breaker로 평균 위치 오름차순으로 합니다.

        :param keyword: 검색 키워드
        :param start_key: 페이지네이션을 위한 시작 키 (n-gram별 {'gram_token': exclusive_start_key, ...})
        :param limit: 각 n-gram 당 DynamoDB 조회 최대 항목 수
        :return: (정렬된 item_id 리스트, 다음 페이지 조회를 위한 end_key 딕셔너리 또는 None)
        """
        # 1. 검색어 전처리 및 2-gram 생성
        processed_value = self._pre_process(keyword)
        query_grams = self._make_n_gram(processed_value, 2)
        total_query = len(query_grams)

        # 2. 각 n-gram에 대해 DynamoDB 조회 및 매칭 정보 수집
        # aggregated_results: { item_id: { "query_indexes": set(), "positions": list[int] } }
        aggregated_results = {}
        result_end_keys = {}

        for q_idx, gram in enumerate(query_grams):
            gram_start_key = start_key.get(gram) if start_key and gram in start_key else None
            items, gram_end_key = self.dynamo_db_api.query_items(
                table_name=self.table_name,
                pk_field='g',
                pk_value=gram,
                sk_condition='gte',
                sk_field='s',
                sk_value=len(processed_value),  # 기준: 전처리된 검색어 길이
                reverse=False,
                limit=limit,
                consistent_read=False,
                start_key=gram_start_key
            )
            if gram_end_key is not None:
                result_end_keys[gram] = gram_end_key
            for item in items:
                item_id = item["i"]
                if item_id not in aggregated_results:
                    aggregated_results[item_id] = {
                        "query_indexes": set(),
                        "positions": []
                    }
                aggregated_results[item_id]["query_indexes"].add(q_idx)
                aggregated_results[item_id]["positions"].append(item["p"])

        # 3. 각 후보 문서에 대해 점수 산출
        scored_items = []
        for item_id, data in aggregated_results.items():
            # Coverage: 매칭된 고유 검색어 토큰 비율
            coverage = len(data["query_indexes"]) / total_query if total_query > 0 else 0

            # Continuity: 매칭된 검색어 토큰 인덱스를 정렬한 후 최대 연속 run 길이 계산
            sorted_indexes = sorted(data["query_indexes"])
            max_run = 1
            current_run = 1
            for i in range(1, len(sorted_indexes)):
                if sorted_indexes[i] == sorted_indexes[i - 1] + 1:
                    current_run += 1
                else:
                    if current_run > max_run:
                        max_run = current_run
                    current_run = 1
            if current_run > max_run:
                max_run = current_run
            continuity = max_run / total_query if total_query > 0 else 0

            # Position Score: 1 / (1 + 평균 매칭 위치)
            avg_position = sum(data["positions"]) / len(data["positions"]) if data["positions"] else 0
            position_score = 1 / (1 + avg_position)

            # 최종 점수: 가중합 (문서 길이 관련 요소는 사용하지 않음)
            alpha, beta, gamma = 0.5, 0.3, 0.2
            final_score = alpha * coverage + beta * continuity + gamma * position_score

            scored_items.append((item_id, final_score, avg_position))

        # 4. 최종 정렬: 최종 점수 내림차순, tie-breaker로 평균 위치 오름차순
        scored_items.sort(key=lambda x: (-x[1], x[2]))
        result_item_ids = [item[0] for item in scored_items]

        # 5. Pagination 처리: end_key 내 Decimal 타입 변환
        result_end_key = result_end_keys if result_end_keys else None
        result_end_key = SearchEngine._convert_decimal(result_end_key) if result_end_key is not None else None

        return result_item_ids, result_end_key

    def search_with_levenshtein(self, keyword: str, start_key: Optional[dict], limit: int = 10) -> tuple[
        list[str], Optional[dict]]:
        """
        기존의 n-gram 기반 후보 추출 후, 각 후보의 원본 텍스트와 검색어 간 Levenshtein 유사도를 적용하여 re-ranking합니다.

        1. 기존 search() 함수를 사용해 후보 item_id와 페이지네이션 정보를 가져옵니다.
        2. 각 후보의 원본 텍스트를 가져와(get_original_text(item_id) 메서드가 존재한다고 가정),
           검색어와의 Levenshtein 유사도를 계산합니다.
        3. 유사도 비율을 기준으로 후보를 내림차순 정렬합니다.

        :param keyword: 검색 키워드
        :param start_key: 페이지네이션을 위한 시작 키
        :param limit: 각 n-gram 당 DynamoDB 조회 최대 항목 수
        :return: (정렬된 item_id 리스트, 다음 페이지 조회를 위한 end_key 딕셔너리 또는 None)
        """
        # 1. n-gram 기반 후보 추출 (기존 search 함수를 재사용)
        candidate_ids, result_end_key = self.search(keyword, start_key, limit)

        # 2. 각 후보의 원본 텍스트 조회 (여기서는 get_original_text(item_id) 메서드가 존재한다고 가정)
        candidate_texts = {}
        for item_id in candidate_ids:
            candidate_texts[item_id] = self.get_original_text(item_id)

        # 3. 검색어와 각 후보의 원본 텍스트 간 Levenshtein 유사도 계산 및 re-ranking
        scores = []
        for item_id, text in candidate_texts.items():
            # 필요한 경우 전처리(예: _pre_process) 적용 가능
            sim = levenshtein_ratio(keyword, text)
            scores.append((item_id, sim))

        # 4. 유사도 기준 내림차순 정렬 (유사도가 높을수록 상위)
        scores.sort(key=lambda x: -x[1])
        sorted_candidate_ids = [item[0] for item in scores]

        return sorted_candidate_ids, result_end_key


if __name__ == '__main__':
    import api

    se = SearchEngine(api.dynamo, 'search-n-gram')

    print("=== 테이블 초기화 ===")
    se.init_table()

    # 1. 테스트 데이터 리스트 (아이템 ID, 텍스트)
    test_data = [
        ("TEST1", "신선한 감자와 고소한 은행"),
        ("TEST2", "감자튀김과 바삭한 감자"),
        ("TEST3", "따끈한 감자탕과 매콤한 국물"),
        ("TEST4", "은행나무 아래서 책 읽기"),
        ("TEST5", "감자와 은행 그리고 사과"),
        ("TEST6", "고소한 두부와 신선한 야채"),
        ("TEST7", "감미로운 음악과 감동적인 영화"),
        ("TEST8", "감기 걸린 감자, 감정이 북받쳐"),
        ("TEST9", "매운 고추로 만든 감기약"),
        ("TEST10", "맛있는 감자전과 김치전"),
        ("TEST11", "은행나무 가로수 길을 걸으며 감미로운 노래"),
        ("TEST12", "고소한 견과류와 함께 먹는 감자칩"),
        ("TEST13", "따뜻한 사과차와 함께하는 감성적인 시간"),
        ("TEST14", "새콤달콤한 감귤과 신선한 과일"),
        ("TEST15", "매운 청양고추와 감칠맛 나는 요리"),
        ("TEST16", "은행에서 대출을 받는 방법"),
        ("TEST17", "감미로운 와인과 함께하는 저녁"),
        ("TEST18", "감동적인 영화와 눈물나는 감정"),
        ("TEST19", "소금과 후추로 간을 맞춘 감자요리"),
        ("TEST20", "바삭한 감자튀김과 매콤한 소스"),
    ]

    # 2. Pagination 테스트를 위한 더미 데이터 개수
    dummy_data_count = 120
    dummy_text = "감자"

    print("\n=== 테스트 데이터 삽입 ===")
    for item_id, text in test_data:
        se.put_data(item_id, text)

    print(f"\n=== Pagination 테스트: DUMMY 데이터 삽입 ({dummy_data_count}개) ===")
    for i in range(dummy_data_count):
        se.put_data(f"DUMMY{i}", dummy_text)

    # 3. 검색어 리스트
    test_queries = ["감", "감자", "은행", "사과", "고추", "감미로운", "고소한", "감동", '사과차와 감성적인']


    # 4. 검색 결과 출력 함수
    def print_search_results(query, first_results, first_end_key):
        print("\n------------------------------")
        print(f"🔍 Search Query: '{query}'")
        print(f"✅ 검색 결과 (유사도 순, 최대 100개): {first_results}")
        print(f"📌 Pagination 정보: {first_end_key}")

        # end_key가 있다면 추가 페이지 조회 테스트
        if first_end_key:
            next_results, next_end_key = se.search(query, first_end_key)
            print(f"🔄 Next Page 검색 결과 (유사도 순, 최대 100개): {next_results}")
            print(f"📌 Next Page Pagination 정보: {next_end_key}")


    print("\n=== 검색 테스트 ===")
    for query in test_queries:
        first_results, first_end_key = se.search(query, None)
        print_search_results(query, first_results, first_end_key)
