import gradio as gr
import tempfile, os, shutil
import requests, json, uuid, time, asyncio
from typing import Any, List
from pydantic import PrivateAttr

# llama-index 관련 임포트
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
from llama_index.core.chat_engine import CondenseQuestionChatEngine
from llama_index.core.llms import CustomLLM, CompletionResponse, CompletionResponseGen, LLMMetadata
from llama_index.core.embeddings import BaseEmbedding

# 설치된 패키지에서 클래스 임포트
from clova_llama_index import ClovaClient, ClovaIndexEmbeddings, ClovaLLM

# import chromadb
from llama_index.core import (
    VectorStoreIndex,
    SimpleDirectoryReader,
    StorageContext,
    SimpleKeywordTableIndex,
    TreeIndex,
    Settings
)
# from llama_index.vector_stores.chroma import ChromaVectorStore
# from IPython.display import Markdown, display


# image_url = "https://www.fpn119.co.kr/imgdata/fpn119_co_kr/202408/2024080548213280.jpg"

# ==========================================
# 4. Gradio 인터페이스 및 RAG 시스템 구현
# ==========================================

# HyperCLOVA API 키 (반드시 실제 API 키로 변경하세요)
clova_api_key = os.environ.get("API_KEY")
print(clova_api_key)
client = ClovaClient(api_key=clova_api_key)

# llama-index 설정: HyperCLOVA LLM 및 임베딩 적용
Settings.llm = ClovaLLM(client)
Settings.embed_model = ClovaIndexEmbeddings(client)
Settings.chunk_size = 4096

chat_engine = None

def chat_with_docs(message, history):
    global chat_engine
    if chat_engine is None:
        return "⚠️ chat_engine을 찾지 못했습니다."
    response = chat_engine.chat(message)
    return response.response

def show_image():
    return "./img/19.jpg"  # 업로드한 파일명
    
with gr.Blocks() as demo:    
    print("with gr.Blocks() as demo:")
    gr.Markdown("## 🤖 HyperCLOVA 기반 문서 챗봇")

    documents = SimpleDirectoryReader('./law', required_exts=[".pdf", ".txt"]).load_data()
    # index = VectorStoreIndex.from_documents(documents)
    # chat_engine = index.as_chat_engine(chat_mode="condense_question", verbose=True)

    nodes = Settings.node_parser.get_nodes_from_documents(documents)

    # 임베딩 함수 초기화 후 컬렉션에 제공
    embed_model = ClovaIndexEmbeddings(client, embed_batch_size=1)
    
    Settings.embed_model = embed_model
    Settings.chunk_size = 1024
    Settings.chunk_overlap = 128

    # # initialize storage context (by default it's in-memory)
    storage_context = StorageContext.from_defaults()
    storage_context.docstore.add_documents(nodes)
        
    index = VectorStoreIndex.from_documents(documents)
    chat_engine = index.as_chat_engine(chat_mode="condense_question", verbose=True)

    embed_model = ClovaIndexEmbeddings(client, embed_batch_size=1)

    print("keyword_index start")
    keyword_index = SimpleKeywordTableIndex(nodes, storage_context=storage_context)
    print("vector_index start")
    vector_index = VectorStoreIndex(nodes, storage_context=storage_context)
    print("tree_index start")
    tree_index = TreeIndex(nodes, storage_context=storage_context)

    from llama_index.core.tools import QueryEngineTool

    vector_query_engine = vector_index.as_query_engine(
        response_mode="tree_summarize", use_async=True, name="vector"
    )
    keyword_query_engine = keyword_index.as_query_engine(
        response_mode="tree_summarize", use_async=True, name="keyword"
    )
    tree_query_engine = tree_index.as_query_engine(
        response_mode="tree_summarize", use_async=True, name="tree"
    )

    vector_tool = QueryEngineTool.from_defaults(
        query_engine=vector_query_engine,
        description=(
            "법률에 대한 전반적인 검색이 필요할 때 활용하세요."
        )
    )
    keyword_tool = QueryEngineTool.from_defaults(
        query_engine=keyword_query_engine,
        description=(
            "키워드로 법률 검색을 진행할때 활용하세요."
        )
    )
    tree_tool = QueryEngineTool.from_defaults(
        query_engine=tree_query_engine,
        description=(
            "법률 전체를 Tree 형태로 만들어서 검색을 할 때 활용합니다."
        )
    )

    from llama_index.core import VectorStoreIndex
    from llama_index.core.objects import ObjectIndex
    
    obj_index = ObjectIndex.from_objects(
        [vector_tool, keyword_tool, tree_tool],
        index_cls=VectorStoreIndex,
    )

    from llama_index.core.query_engine import ToolRetrieverRouterQueryEngine
    query_engine = ToolRetrieverRouterQueryEngine(obj_index.as_retriever())

    # response = query_engine.query("소방시설업을 등록하려는 사람은 어떤 절차를 따라야 하나요? 관련된 법령 조항을 함께 표시해 주세요. 답변을 이해하기 쉽게 최대한 상세히 작성해 주세요")
    # print(response)

    # with gr.Row():
    #     file_input = gr.File(file_types=[".pdf", ".txt"], file_count="multiple", label="📁 문서 업로드")
    #     upload_btn = gr.Button("📌 문서 분석 시작")
    # status_box = gr.Textbox(label="상태", interactive=False)

    # 최신 Gradio 형식에 맞게 type="messages" 설정
    chatbot = gr.ChatInterface(fn=chat_with_docs, title="📚 문서 기반 챗봇", type="messages") 

    # upload_btn.click(fn=upload_and_prepare_chat, inputs=[file_input], outputs=[status_box])

demo.launch()