import pandas as pd
import numpy as np
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain.schema import Document
from rank_bm25 import BM25Okapi
from kiwipiepy import Kiwi
from typing import List
import gradio as gr

class ProductSearchSystem:
    def __init__(self, 
                 model_name: str = "snunlp/KR-SBERT-V40K-klueNLI-augSTS",
                 bm25_weight: float = 0.3,
                 vector_weight: float = 0.7):
        """검색 시스템 초기화"""
        self.embeddings = HuggingFaceEmbeddings(
            model_name=model_name,
            model_kwargs={'device': 'cpu'},
            encode_kwargs={'normalize_embeddings': True}
        )
        self.bm25_weight = bm25_weight
        self.vector_weight = vector_weight
        self.vector_store = None
        self.bm25 = None
        self.documents = []
        self.df = None
        # Kiwi 토크나이저 초기화
        self.kiwi = Kiwi()
        
    def _tokenize_text(self, text: str) -> List[str]:
        """Kiwi를 사용한 텍스트 토크나이징"""
        # 형태소 분석 수행
        tokens = self.kiwi.tokenize(text)
        # 명사, 동사, 형용사만 추출
        pos_tags = ['NNG', 'NNP', 'VV', 'VA', 'SL']  # 일반명사, 고유명사, 동사, 형용사
        return [token.form for token in tokens if token.tag in pos_tags]  # pos를 tag로 변경
    
    def load_sample_data(self):
        """샘플 데이터 로드"""
        self.df = pd.read_csv("sample_data.csv")
        self._preprocess_data()
        self._create_search_index()
        return True

    def _preprocess_data(self):
        """데이터 전처리"""
        # 빈 값 처리
        self.df['category'] = self.df['category'].fillna('미분류')
        
        # 특수 문자 처리
        self.df['company_info'] = self.df['company_info'].fillna('')
        self.df['company_info'] = self.df['company_info'].str.replace('_x000D_', '\n')
        self.df['description'] = self.df['description'].fillna('')
        self.df['description'] = self.df['description'].str.replace('_x000D_', '\n')
        
        # 불필요한 공백 제거
        for col in self.df.columns:
            if self.df[col].dtype == 'object':
                self.df[col] = self.df[col].str.strip()
    
    def _create_search_index(self):
        """검색 인덱스 생성"""
        self.documents = []
        tokenized_documents = []  # BM25용 토큰화된 문서
        
        for _, row in self.df.iterrows():
            content = f"{row['company_name']} {row['category']} {row['company_info']} {row['product_name']} {row['description']}"
            # Kiwi 토크나이저를 사용한 토큰화
            tokenized_doc = self._tokenize_text(content)
            tokenized_documents.append(tokenized_doc)
            
            self.documents.append(
                Document(
                    page_content=content,
                    metadata={
                        'company_name': row['company_name'],
                        'category': row['category'],
                        'company_info': row['company_info'],
                        'product_name': row['product_name'],
                        'description': row['description']
                    }
                )
            )
        
        # BM25 인덱스 생성
        self.bm25 = BM25Okapi(tokenized_documents)
        
        # 벡터 스토어 생성
        self.vector_store = FAISS.from_documents(self.documents, self.embeddings)
    
    def search(self, query: str, top_k: int = 3) -> List[dict]:
        """검색 실행"""
        if not query.strip():
            return []
        
        # BM25 검색 - Kiwi 토크나이저 사용
        tokenized_query = self._tokenize_text(query)
        bm25_scores = self.bm25.get_scores(tokenized_query)
        
        # 벡터 검색
        query_embedding = self.embeddings.embed_query(query)
        vector_docs_and_scores = self.vector_store.similarity_search_with_score(query, k=len(self.documents))
        
        # 결과 통합 및 점수 계산
        results = []
        seen_products = set()
        
        # 점수 정규화를 위한 최대값
        max_bm25 = max(bm25_scores) if len(bm25_scores) > 0 else 1
        max_vector = max(score for _, score in vector_docs_and_scores) if vector_docs_and_scores else 1
        
        for i, doc in enumerate(self.documents):
            # 정규화된 점수 계산
            bm25_score = bm25_scores[i] / max_bm25 if max_bm25 > 0 else 0
            vector_score = None
            
            # 해당 문서의 벡터 점수 찾기
            for vec_doc, vec_score in vector_docs_and_scores:
                if vec_doc.page_content == doc.page_content:
                    vector_score = (1 - (vec_score / max_vector)) if max_vector > 0 else 0
                    break
            
            if vector_score is not None:
                # 최종 점수 계산
                final_score = (self.bm25_weight * bm25_score) + (self.vector_weight * vector_score)
                
                product_key = f"{doc.metadata['company_name']}-{doc.metadata['product_name']}"
                if product_key not in seen_products:
                    results.append({
                        'company_name': doc.metadata['company_name'],
                        'category': doc.metadata['category'],
                        'company_info': doc.metadata['company_info'],
                        'product_name': doc.metadata['product_name'],
                        'description': doc.metadata['description'],
                        'bm25_score': round(bm25_score, 3),
                        'vector_score': round(vector_score, 3),
                        'final_score': round(final_score, 3)
                    })
                    seen_products.add(product_key)
        
        # 최종 점수로 정렬
        results.sort(key=lambda x: x['final_score'], reverse=True)
        return results[:top_k]

def create_gradio_interface():
    """Gradio 인터페이스 생성"""
    # 검색 시스템 초기화 및 샘플 데이터 로드
    search_system = ProductSearchSystem()
    search_system.load_sample_data()
    
    def search_products(query: str, 
                       top_k: int, 
                       bm25_weight: float) -> tuple:
        """검색 실행 및 결과 포매팅"""
        # 가중치 업데이트
        search_system.bm25_weight = bm25_weight
        search_system.vector_weight = 1 - bm25_weight
        
        # 검색 실행
        results = search_system.search(query, top_k=top_k)
        
        # 결과를 표 형식으로 변환
        if results:
            # 표시할 열 순서 지정
            columns_order = ['company_name', 'category', 'company_info', 'product_name', 'bm25_score', 'vector_score', 'final_score', 'description']
            df_results = pd.DataFrame(results)[columns_order]
            
            # 열 이름 한글화
            df_results.columns = ['회사명', '카테고리', '회사 설명', '제품명', '키워드 점수', '벡터 점수', '최종 점수', '설명']
            
            html_table = df_results.to_html(
                classes=['table', 'table-striped'], 
                escape=False, 
                index=False,
                float_format=lambda x: '{:.3f}'.format(x)  # 소수점 3자리까지 표시
            )
        else:
            html_table = "<p>검색 결과가 없습니다.</p>"
        
        # 상세 결과 텍스트 생성
        detailed_results = []
        for i, result in enumerate(results, 1):
            detailed_results.append(f"""
            === 검색결과 #{i} ===
            회사명: {result['company_name']}
            카테고리: {result['category']}
            회사 설명: {result['company_info']}
            제품명: {result['product_name']}
            키워드 점수: {result['bm25_score']:.3f}
            벡터 점수: {result['vector_score']:.3f}
            최종 점수: {result['final_score']:.3f}
            설명: {result['description']}
            """)
        
        detailed_text = "\n".join(detailed_results) if detailed_results else "검색 결과가 없습니다."
        
        return html_table, detailed_text
    
    # Gradio 인터페이스 정의
    with gr.Blocks(css="footer {visibility: hidden}") as demo:
        gr.Markdown("""
        # 🔍 코엑스 부스 추천 시스템
        하이브리드 방식을 이용한 기업 및 제품 검색/추천 시스템입니다.
        """)
        
        with gr.Row():
            with gr.Column(scale=4):
                query_input = gr.Textbox(
                    label="검색어를 입력하세요",
                    placeholder="예: AI 기술 회사, 센서, 자동화 등",
                )
            with gr.Column(scale=1):
                top_k = gr.Slider(
                    minimum=1,
                    maximum=10,
                    value=3,
                    step=1,
                    label="검색 결과 수",
                )
        
        with gr.Row():
            bm25_weight = gr.Slider(
                minimum=0.0,
                maximum=1.0,
                value=0.3,
                step=0.1,
                label="키워드 검색 가중치",
            )
        
        with gr.Row():
            search_button = gr.Button("검색", variant="primary")
        
        with gr.Row():
            with gr.Column():
                results_table = gr.HTML(label="검색 결과 테이블")
            with gr.Column():
                results_text = gr.Textbox(
                    label="상세 결과",
                    show_label=True,
                    interactive=False,
                    lines=10
                )
        
        # 이벤트 핸들러 연결
        search_button.click(
            fn=search_products,
            inputs=[query_input, top_k, bm25_weight],
            outputs=[results_table, results_text],
        )
            
        gr.Markdown("""
        ### 사용 방법
        1. 검색어 입력: 찾고자 하는 기업, 제품, 기술 등의 키워드를 입력하세요
        2. 검색 결과 수 조정: 원하는 결과 수를 선택하세요
        3. 가중치 조정: 키워드 매칭과 의미적 유사도 간의 가중치를 조절하세요

        ### 점수 설명
        - 키워드 점수: Kiwi 토크나이저를 사용한 키워드 기반 매칭 점수 (0~1)
        - 벡터 점수: 의미적 유사도 점수 (0~1)
        - 최종 점수: 키워드 점수와 벡터 점수의 가중 평균
        """)
    
    return demo

def main():
    demo = create_gradio_interface()
    demo.launch(share=True)

if __name__ == "__main__":
    main()




# TODO
# OCR 딥러닝 vs OCR 처리
# 토크나이저 처리 결과 테스트
# 품사 태깅 결과 확인