import pandas as pd
import itertools
from sentence_transformers import SentenceTransformer
from pinecone import Pinecone

# --- データ前処理 ---

# CSVファイルからルール情報読み込み
df = pd.read_csv("abcd.csv")
filtered_df = df[["カテゴリー", "Q", "A"]]

rule_data_list = []
for _, row in filtered_df.iterrows():
    category = row["カテゴリー"]
    question = str(row["Q"]).strip()
    answer = str(row["A"]).strip()
    combined = f"以下は、「{category}」に関するルール情報です。{question} {answer}"
    rule_data_list.append(combined)

# Slackデータの整形
#Slackチャンネルから取得した情報例
slack_messages = [
    "○○ソフトウェアをインストールしたいです。",
]
channel_id = "C12342391"
slack_data_list = [
    f"以下は、Slackチャンネル「{channel_id}」の情報です。{msg.strip()}"
    for msg in slack_messages
]

# --- ベクトルDB登録処理 ---

# Pinecone設定
#プロジェクトのAPIキー
pc = Pinecone(api_key='******') #先ほど作成したAPIキーの情報を設定
index = pc.Index(host="https://****.pinecone.io")#先ほど作成したindexのHOST URLを設定

# HuggingFaceからモデル読み込み（GluCoSE）
model = SentenceTransformer("pkshatech/GLuCoSE-base-ja-v2")

# ベクトル生成
def create_embeddings(sentence):
    return model.encode(f"passage: {sentence}").tolist()

# レコード生成
def generate_records(data_list, prefix, start_id):
    is_f_flag = prefix == "F"
    return [
        {
            "id": f"{prefix}{str(start_id + i)}",
            "values": create_embeddings(data),
            "metadata": {"doc": data, "is_f": is_f_flag}
        }
        for i, data in enumerate(data_list)
    ]

# チャンク分割
def chunks(iterable, batch_size):
    it = iter(iterable)
    chunk = tuple(itertools.islice(it, batch_size))
    while chunk:
        yield chunk
        chunk = tuple(itertools.islice(it, batch_size))

# Upsert実行
def do_upsert(record_list, batch_size):
    for chunk in chunks(record_list, batch_size=batch_size):
        index.upsert(vectors=chunk)
    print("UPSERTED:", len(record_list))

# 実行
rule_records = generate_records(rule_data_list, prefix="F", start_id=1)
slack_records = generate_records(slack_data_list, prefix="S", start_id=1)

do_upsert(rule_records, batch_size=100)
do_upsert(slack_records, batch_size=100)

from sentence_transformers import SentenceTransformer
import os
from pinecone import Pinecone
import logging

logger = logging.getLogger()
logger.setLevel(logging.INFO)

#プロジェクトのAPIキー
pc = Pinecone(api_key='******') #先ほど作成したAPIキーの情報を設定
index = pc.Index(host="https://****.pinecone.io")#先ほど作成したindexのHOST URLを設定

# --- Load GLuCoSE model ---
model = SentenceTransformer("pkshatech/GLuCoSE-base-ja-v2")

def exec_rag(question, similarity):  # similarity: -1 ~ 1
    try:
        question_text = f"query: {question}"
        q_embeddings = model.encode(question_text).tolist()
        print("Embeddings created.")
        print(q_embeddings)

        print("Querying Pinecone for similarity search...")
        # For File IDs (is_f=True)
        search_results_f = index.query(vector=q_embeddings, top_k=2,filter={"is_f": True}, include_values=False, include_metadata=True)
        # For Slack IDs (is_f not True OR missing)
        search_results_s = index.query(
            vector=q_embeddings,
            top_k=2,
            filter={
                "$or": [
                    {"is_f": {"$ne": True}},  # is_f exists but is not True
                    {"is_f": {"$exists": False}}  # is_f not set (old Slack vectors)
                ]
            },
            include_values=False,
            include_metadata=True
        )
        # Filter search results based on similarity
        context_candidates = []
        work_list = []

        filtered_f = [
            [val['id'], val['score'], val['metadata']['doc']]
            for val in search_results_f['matches']
            if val['score'] >= similarity
        ]
        filtered_s = [
            [val['id'], val['score'], val['metadata']['doc']]
            for val in search_results_s['matches']
            if val['score'] >= similarity
        ]

        # Sort only by ID (descending)
        filtered_f.sort(key=lambda x: int(x[0][1:]), reverse=True)
        filtered_s.sort(key=lambda x: int(x[0][1:]), reverse=True)

        # Create Context 検索結果 with F first, then S
        context_candidates = filtered_f + filtered_s
        print(context_candidates)

    except Exception as e:
        print(f"Error in exec_rag: {str(e)}")
        raise


# 質問と検索結果
question = "サクラエディタをインストールしたいです。"
similarity = 0.65
exec_rag(question, similarity)