package main

import (
	"context"
	"fmt"
	"strings"
	"sync"
	"time"
)

// EntityNode 表示知识图谱中的实体节点
type EntityNode struct {
	ID   string
	Name string
	Type string // 实体类型：公司、人物、产品等
}

// RelationEdge 表示实体之间的关系
type RelationEdge struct {
	SourceID string
	TargetID string
	Relation string
	Weight   float64 // 关系权重
}

// KnowledgeGraph 知识图谱结构
type KnowledgeGraph struct {
	entities map[string]*EntityNode
	edges    []*RelationEdge
	adjList  map[string][]*RelationEdge // 邻接表加速检索
	mu       sync.RWMutex
}

// NewKnowledgeGraph 创建知识图谱实例
func NewKnowledgeGraph() *KnowledgeGraph {
	return &KnowledgeGraph{
		entities: make(map[string]*EntityNode),
		edges:    make([]*RelationEdge, 0),
		adjList:  make(map[string][]*RelationEdge),
	}
}

// AddEntity 添加实体节点
func (kg *KnowledgeGraph) AddEntity(entity *EntityNode) {
	kg.mu.Lock()
	defer kg.mu.Unlock()
	kg.entities[entity.ID] = entity
}

// AddRelation 添加关系边
func (kg *KnowledgeGraph) AddRelation(edge *RelationEdge) {
	kg.mu.Lock()
	defer kg.mu.Unlock()
	kg.edges = append(kg.edges, edge)
	kg.adjList[edge.SourceID] = append(kg.adjList[edge.SourceID], edge)
}

// EntityLinker 实体链接器
type EntityLinker struct {
	kg        *KnowledgeGraph
	synonyms  map[string]string // 同义词映射：用户输入 -> 标准实体名
	threshold float64           // 链接置信度阈值
}

// NewEntityLinker 创建实体链接器
func NewEntityLinker(kg *KnowledgeGraph) *EntityLinker {
	return &EntityLinker{
		kg:        kg,
		synonyms:  make(map[string]string),
		threshold: 0.6,
	}
}

// AddSynonym 添加同义词映射
func (el *EntityLinker) AddSynonym(input, standard string) {
	el.synonyms[input] = standard
}

// LinkEntity 将文本中的实体提及链接到知识图谱实体
func (el *EntityLinker) LinkEntity(ctx context.Context, mention string) ([]*EntityNode, float64) {
	// 步骤1：检查同义词映射
	if standard, ok := el.synonyms[mention]; ok {
		mention = standard
	}

	// 步骤2：在知识图谱中查找匹配实体
	var candidates []*EntityNode
	for _, entity := range el.kg.entities {
		// 使用简单字符串匹配，生产环境应使用模糊匹配或嵌入相似度
		if strings.Contains(entity.Name, mention) || strings.Contains(mention, entity.Name) {
			candidates = append(candidates, entity)
		}
	}

	// 步骤3：计算置信度
	if len(candidates) == 0 {
		return nil, 0.0
	}

	// 简单置信度计算：匹配度最高的实体
	confidence := 0.8 // 假设匹配成功
	return candidates, confidence
}

// GraphRetriever 图结构检索器
type GraphRetriever struct {
	kg *KnowledgeGraph
}

// NewGraphRetriever 创建图检索器
func NewGraphRetriever(kg *KnowledgeGraph) *GraphRetriever {
	return &GraphRetriever{kg: kg}
}

// RetrieveByPath 根据路径模式检索信息
// pathPattern 示例：[{"relation":"竞争对手"}, {"relation":"销量","direction":"out"}]
func (gr *GraphRetriever) RetrieveByPath(ctx context.Context, startEntityID string, pathPattern []string) ([]map[string]interface{}, error) {
	results := make([]map[string]interface{}, 0)
	visited := make(map[string]bool)

	// BFS遍历路径
	queue := []struct {
		nodeID string
		path   []string
		info   map[string]interface{}
	}{{nodeID: startEntityID, path: []string{}, info: make(map[string]interface{})}}

	for len(queue) > 0 {
		current := queue[0]
		queue = queue[1:]

		if visited[current.nodeID] {
			continue
		}
		visited[current.nodeID] = true

		// 检查当前节点是否满足路径模式
		if len(current.path) == len(pathPattern) {
			results = append(results, current.info)
			continue
		}

		// 获取当前节点的邻接边
		edges := gr.kg.adjList[current.nodeID]
		nextRelation := pathPattern[len(current.path)]

		for _, edge := range edges {
			if edge.Relation == nextRelation {
				newInfo := make(map[string]interface{})
				for k, v := range current.info {
					newInfo[k] = v
				}
				// 记录检索到的三元组
				newInfo[edge.Relation] = gr.kg.entities[edge.TargetID].Name

				queue = append(queue, struct {
					nodeID string
					path   []string
					info   map[string]interface{}
				}{
					nodeID: edge.TargetID,
					path:   append(current.path, edge.Relation),
					info:   newInfo,
				})
			}
		}
	}

	return results, nil
}

// FusionEngine 融合推理引擎
type FusionEngine struct {
	maxTokens int // 最大融合上下文长度
}

// NewFusionEngine 创建融合引擎
func NewFusionEngine(maxTokens int) *FusionEngine {
	return &FusionEngine{maxTokens: maxTokens}
}

// FusionContext 融合后的上下文信息
type FusionContext struct {
	GraphTriples []string // 图结构三元组文本化
	Entities     []string // 涉及的实体列表
	Relations    []string // 涉及的关系列表
	Summary      string   // 融合摘要
}

// Fuse 执行信息融合
func (fe *FusionEngine) Fuse(ctx context.Context, query string, graphResults []map[string]interface{}) *FusionContext {
	fc := &FusionContext{
		GraphTriples: make([]string, 0),
		Entities:     make([]string, 0),
		Relations:    make([]string, 0),
	}

	// 将图检索结果转化为自然语言三元组
	for _, result := range graphResults {
		for relation, value := range result {
			triple := fmt.Sprintf("(%s, %s, %v)", query, relation, value)
			fc.GraphTriples = append(fc.GraphTriples, triple)
			fc.Relations = append(fc.Relations, relation)
		}
	}

	// 生成融合摘要
	if len(fc.GraphTriples) > 0 {
		fc.Summary = fmt.Sprintf("根据知识图谱检索，查询'%s'涉及以下事实：%s",
			query, strings.Join(fc.GraphTriples, "; "))
	} else {
		fc.Summary = fmt.Sprintf("未从知识图谱中找到与'%s'直接相关的事实", query)
	}

	return fc
}

// RAGSystem 完整的RAG系统
type RAGSystem struct {
	linker    *EntityLinker
	retriever *GraphRetriever
	fusion    *FusionEngine
	llm       func(string) string // 模拟LLM生成函数
}

// NewRAGSystem 创建RAG系统实例
func NewRAGSystem(kg *KnowledgeGraph, llm func(string) string) *RAGSystem {
	return &RAGSystem{
		linker:    NewEntityLinker(kg),
		retriever: NewGraphRetriever(kg),
		fusion:    NewFusionEngine(4096),
		llm:       llm,
	}
}

// Answer 回答用户查询
func (rs *RAGSystem) Answer(ctx context.Context, query string) string {
	// 步骤1：实体链接
	entities, confidence := rs.linker.LinkEntity(ctx, query)
	if confidence < rs.linker.threshold {
		return "无法准确理解您的查询中的实体信息，请提供更具体的描述。"
	}

	// 步骤2：图结构检索
	var allResults []map[string]interface{}
	for _, entity := range entities {
		// 根据查询意图构建路径模式，这里简单使用固定模式
		pathPattern := []string{"竞争对手", "销量"}
		results, err := rs.retriever.RetrieveByPath(ctx, entity.ID, pathPattern)
		if err == nil {
			allResults = append(allResults, results...)
		}
	}

	// 步骤3：信息融合
	fusionContext := rs.fusion.Fuse(ctx, query, allResults)

	// 步骤4：生成最终答案
	prompt := fmt.Sprintf("请根据以下知识图谱信息回答用户问题。\n知识图谱信息：%s\n用户问题：%s",
		fusionContext.Summary, query)

	return rs.llm(prompt)
}

func main() {
	// 初始化知识图谱
	kg := NewKnowledgeGraph()

	// 添加实体
	kg.AddEntity(&EntityNode{ID: "1", Name: "特斯拉", Type: "公司"})
	kg.AddEntity(&EntityNode{ID: "2", Name: "比亚迪", Type: "公司"})
	kg.AddEntity(&EntityNode{ID: "3", Name: "埃隆·马斯克", Type: "人物"})
	kg.AddEntity(&EntityNode{ID: "4", Name: "王传福", Type: "人物"})

	// 添加关系
	kg.AddRelation(&RelationEdge{SourceID: "1", TargetID: "2", Relation: "竞争对手", Weight: 0.9})
	kg.AddRelation(&RelationEdge{SourceID: "1", TargetID: "3", Relation: "CEO", Weight: 1.0})
	kg.AddRelation(&RelationEdge{SourceID: "2", TargetID: "4", Relation: "CEO", Weight: 1.0})
	kg.AddRelation(&RelationEdge{SourceID: "1", TargetID: "2", Relation: "销量比较", Weight: 0.8})

	// 模拟LLM生成函数
	llm := func(prompt string) string {
		return fmt.Sprintf("基于知识图谱分析，特斯拉的竞争对手包括比亚迪。特斯拉的CEO是埃隆·马斯克，比亚迪的CEO是王传福。两者在新能源汽车市场存在直接竞争关系。")
	}

	// 创建RAG系统
	rag := NewRAGSystem(kg, llm)

	// 测试查询
	ctx := context.Background()
	answer := rag.Answer(ctx, "特斯拉的竞争对手是谁")
	fmt.Println("系统回答:", answer)
}