import os
import sys

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import json
from typing import List, Dict, Tuple
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import glob
import sqlite3
import re
from datetime import datetime
import random
from src.llm_models.utils_model import LLMRequest
from src.config.config import global_config


def clean_group_name(name: str) -> str:
    """清理群组名称，只保留中文和英文字符"""
    cleaned = re.sub(r"[^\u4e00-\u9fa5a-zA-Z]", "", name)
    if not cleaned:
        cleaned = datetime.now().strftime("%Y%m%d")
    return cleaned


def get_group_name(stream_id: str) -> str:
    """从数据库中获取群组名称"""
    conn = sqlite3.connect("data/maibot.db")
    cursor = conn.cursor()

    cursor.execute(
        """
        SELECT group_name, user_nickname, platform 
        FROM chat_streams 
        WHERE stream_id = ?
    """,
        (stream_id,),
    )

    result = cursor.fetchone()
    conn.close()

    if result:
        group_name, user_nickname, platform = result
        if group_name:
            return clean_group_name(group_name)
        if user_nickname:
            return clean_group_name(user_nickname)
        if platform:
            return clean_group_name(f"{platform}{stream_id[:8]}")
    return stream_id


def load_expressions(chat_id: str) -> List[Dict]:
    """加载指定群聊的表达方式"""
    style_file = os.path.join("data", "expression", "learnt_style", str(chat_id), "expressions.json")

    style_exprs = []

    if os.path.exists(style_file):
        with open(style_file, "r", encoding="utf-8") as f:
            style_exprs = json.load(f)

    # 如果表达方式超过10个，随机选择10个
    if len(style_exprs) > 50:
        style_exprs = random.sample(style_exprs, 50)
        print(f"\n从 {len(style_exprs)} 个表达方式中随机选择了 10 个进行匹配")

    return style_exprs


def find_similar_expressions_tfidf(
    input_text: str, expressions: List[Dict], mode: str = "both", top_k: int = 10
) -> List[Tuple[str, str, float]]:
    """使用TF-IDF方法找出与输入文本最相似的top_k个表达方式"""
    if not expressions:
        return []

    # 准备文本数据
    if mode == "style":
        texts = [expr["style"] for expr in expressions]
    elif mode == "situation":
        texts = [expr["situation"] for expr in expressions]
    else:  # both
        texts = [f"{expr['situation']} {expr['style']}" for expr in expressions]

    texts.append(input_text)  # 添加输入文本

    # 使用TF-IDF向量化
    vectorizer = TfidfVectorizer()
    tfidf_matrix = vectorizer.fit_transform(texts)

    # 计算余弦相似度
    similarity_matrix = cosine_similarity(tfidf_matrix)

    # 获取输入文本的相似度分数（最后一行）
    scores = similarity_matrix[-1][:-1]  # 排除与自身的相似度

    # 获取top_k的索引
    top_indices = np.argsort(scores)[::-1][:top_k]

    # 获取相似表达
    similar_exprs = []
    for idx in top_indices:
        if scores[idx] > 0:  # 只保留有相似度的
            similar_exprs.append((expressions[idx]["style"], expressions[idx]["situation"], scores[idx]))

    return similar_exprs


async def find_similar_expressions_embedding(
    input_text: str, expressions: List[Dict], mode: str = "both", top_k: int = 5
) -> List[Tuple[str, str, float]]:
    """使用嵌入模型找出与输入文本最相似的top_k个表达方式"""
    if not expressions:
        return []

    # 准备文本数据
    if mode == "style":
        texts = [expr["style"] for expr in expressions]
    elif mode == "situation":
        texts = [expr["situation"] for expr in expressions]
    else:  # both
        texts = [f"{expr['situation']} {expr['style']}" for expr in expressions]

    # 获取嵌入向量
    llm_request = LLMRequest(global_config.model.embedding)
    text_embeddings = []
    for text in texts:
        embedding = await llm_request.get_embedding(text)
        if embedding:
            text_embeddings.append(embedding)

    input_embedding = await llm_request.get_embedding(input_text)
    if not input_embedding or not text_embeddings:
        return []

    # 计算余弦相似度
    text_embeddings = np.array(text_embeddings)
    similarities = np.dot(text_embeddings, input_embedding) / (
        np.linalg.norm(text_embeddings, axis=1) * np.linalg.norm(input_embedding)
    )

    # 获取top_k的索引
    top_indices = np.argsort(similarities)[::-1][:top_k]

    # 获取相似表达
    similar_exprs = []
    for idx in top_indices:
        if similarities[idx] > 0:  # 只保留有相似度的
            similar_exprs.append((expressions[idx]["style"], expressions[idx]["situation"], similarities[idx]))

    return similar_exprs


async def main():
    # 获取所有群聊ID
    style_dirs = glob.glob(os.path.join("data", "expression", "learnt_style", "*"))
    chat_ids = [os.path.basename(d) for d in style_dirs]

    if not chat_ids:
        print("没有找到任何群聊的表达方式数据")
        return

    print("可用的群聊:")
    for i, chat_id in enumerate(chat_ids, 1):
        group_name = get_group_name(chat_id)
        print(f"{i}. {group_name}")

    while True:
        try:
            choice = int(input("\n请选择要分析的群聊编号 (输入0退出): "))
            if choice == 0:
                break
            if 1 <= choice <= len(chat_ids):
                chat_id = chat_ids[choice - 1]
                break
            print("无效的选择，请重试")
        except ValueError:
            print("请输入有效的数字")

    if choice == 0:
        return

    # 加载表达方式
    style_exprs = load_expressions(chat_id)

    group_name = get_group_name(chat_id)
    print(f"\n已选择群聊：{group_name}")

    # 选择匹配模式
    print("\n请选择匹配模式：")
    print("1. 匹配表达方式")
    print("2. 匹配情景")
    print("3. 两者都考虑")

    while True:
        try:
            mode_choice = int(input("\n请选择匹配模式 (1-3): "))
            if 1 <= mode_choice <= 3:
                break
            print("无效的选择，请重试")
        except ValueError:
            print("请输入有效的数字")

    mode_map = {1: "style", 2: "situation", 3: "both"}
    mode = mode_map[mode_choice]

    # 选择匹配方法
    print("\n请选择匹配方法：")
    print("1. TF-IDF方法")
    print("2. 嵌入模型方法")

    while True:
        try:
            method_choice = int(input("\n请选择匹配方法 (1-2): "))
            if 1 <= method_choice <= 2:
                break
            print("无效的选择，请重试")
        except ValueError:
            print("请输入有效的数字")

    while True:
        input_text = input("\n请输入要匹配的文本（输入q退出）: ")
        if input_text.lower() == "q":
            break

        if not input_text.strip():
            continue

        if method_choice == 1:
            similar_exprs = find_similar_expressions_tfidf(input_text, style_exprs, mode)
        else:
            similar_exprs = await find_similar_expressions_embedding(input_text, style_exprs, mode)

        if similar_exprs:
            print("\n找到以下相似表达：")
            for style, situation, score in similar_exprs:
                print(f"\n\033[33m表达方式：{style}\033[0m")
                print(f"\033[32m对应情景：{situation}\033[0m")
                print(f"相似度：{score:.3f}")
                print("-" * 20)
        else:
            print("\n没有找到相似的表达方式")


if __name__ == "__main__":
    import asyncio

    asyncio.run(main())
