# -*- coding: utf-8 -*-
"""
@file : faiss_functions.py
@author: jwkim
@license: A2D2

@section Modify History
- 2025-06-27 오후 2:05 hsj100 base

"""

"""
Basic logger: custom color logger
"""
try:
    from custom_logger.custom_log import custom_logger as log
except ImportError:
    import logging as log; log.basicConfig(level=log.DEBUG)


"""
Package: 3rd party
"""

"""
Package: standard
"""
import random
import pathlib
import os
import sys
import traceback

"""
Package: custom
"""
# NOTE(hsj100): 코드 단독 실행을 위한 기본 경로 지정 (일반적으로는 본 코드 보다는 모듈 실행 방법을 권장)

if __name__ == '__main__':
    # 현재 스크립트의 디렉토리를 기준으로 부모 디렉토리의 경로를 가져옴
    parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
    if parent_dir not in sys.path:
        sys.path.append(parent_dir)

from custom_apps.FEATURE_VECTOR_SIMILARITY_FAISS.report_utils import *
from custom_apps.FEATURE_VECTOR_SIMILARITY_FAISS.const import *
from custom_apps.FEATURE_VECTOR_SIMILARITY_FAISS.faiss_similarity_search import VectorSimilarity
from custom_apps.FEATURE_VECTOR_SIMILARITY_FAISS.report_utils import *
"""
Definition
"""


"""
Implementation
"""

def get_models(index_type, model_type):
    from vector_rest.app import models as M
    
    model = None
    
    # model 선택
    if index_type == M.VitIndexType.l2:
        # l2
        if model_type == M.VitModelType.l14:
            model = FEMUsageInfo.openaiclip_vit_l14_l2
        elif model_type == M.VitModelType.b32:
            model = FEMUsageInfo.openaiclip_vit_b32_l2
        elif model_type == M.VitModelType.b16:
            model = FEMUsageInfo.openaiclip_vit_b16_l2
        elif model_type == M.VitModelType.l14_336:
            model = FEMUsageInfo.openaiclip_vit_l14_336_l2
        else:
            raise Exception(f"model_type {model_type} is invalid")
    else:
        #cosine
        if model_type == M.VitModelType.l14:
            model = FEMUsageInfo.openaiclip_vit_l14_cos
        elif model_type == M.VitModelType.b32:
            model = FEMUsageInfo.openaiclip_vit_b32_cos
        elif model_type == M.VitModelType.b16:
            model = FEMUsageInfo.openaiclip_vit_b16_cos
        elif model_type == M.VitModelType.l14_336:
            model = FEMUsageInfo.openaiclip_vit_l14_336_cos
        else:
            # NOTE(JWKIM) : 현재 l14만 지원 FEMUsageInfo 에 추가해야함
            raise Exception(f"model_type {model_type} is invalid")
        
    if model:
        return model
    else:
        raise Exception(f"model is not selected")
    
    
def save_report_image(report_info, report_image_path):
    """
    이미지 유사도 검색 결과를 리포트로 생성
    """
    table_setting(report_info, report_image_path)


def get_clip_info(model, query_image_path, top_k=4):
    """
    이미지 유사도 검색
    """
    index_file_path = os.path.join(VECTOR_IMG_LIST_PATH,f'{model.name}_{os.path.basename(model.value[1].trained_model)}_{model.value[1].index_type}_{DEFAULT_INDEX_NAME_SUFFIX}')
    txt_file_path = os.path.join(VECTOR_IMG_LIST_PATH,f'{model.name}_{os.path.basename(model.value[1].trained_model)}_{model.value[1].index_type}.txt')
    vector_model = VectorSimilarity(model_info=model, 
                            index_file_path=index_file_path,
                            txt_file_path=txt_file_path)

    # vector_model.create_feature_extraction_model(FEMUsageInfo.openaiclipvit)

    vector_model.save_index_from_files(images_path=INDEX_IMAGES_PATH,
                            save_index_dir=FAISS_VECTOR_PATH,
                            save_txt_dir=VECTOR_IMG_LIST_PATH,
                            index_type=model.value[1].index_type)
                                
    inference_times, result_img_paths, result_percents = vector_model.query_faiss(query_image_path, top_k=top_k)

    report_info = ReportInfo(
        feature_extraction_model=ReportInfoConst.feature_extraction_model,
        model_architecture=ReportInfoConst.model_architecture,
        train_model=os.path.basename(model.value[1].trained_model),
        feature_extraction_time=inference_times,
        query_image_path=query_image_path,
        result_image_paths=result_img_paths,
        result_percents=result_percents,
        similarity_type=model.value[1].index_type
    )
    
    return report_info
    
    
