package llm

import (
	"context"
	"encoding/json"
	"fmt"
	"knowledge-graph-backend/internal/model"
	"log"

	"github.com/tmc/langchaingo/llms"
)

// KnowledgeGraphToolService 定义 LLM 客户端需要的图谱操作接口
// service.Neo4jService 会隐式实现该接口，避免循环依赖
type KnowledgeGraphToolService interface {
	SearchNodes(query string) []model.Node
	GetNeighbors(nodeID string) (model.NeighborResponse, bool)
	GetNodeByID(id string) (model.Node, bool)
	CreateNode(req model.CreateNodeRequest) (model.Node, error)
	CreateEdge(req model.CreateEdgeRequest) (model.Edge, error)
	DeleteNode(id string) error
	DeleteEdge(id string) error
}

// Agent 封装了带有工具调用和上下文管理能力的智能体
type Agent struct {
	llm      llms.Model
	graphSvc KnowledgeGraphToolService
}

// NewAgent 创建一个新的 Agent 实例
func NewAgent(llm llms.Model, graphSvc KnowledgeGraphToolService) *Agent {
	return &Agent{
		llm:      llm,
		graphSvc: graphSvc,
	}
}

// getKnowledgeGraphTools 声明 LLM 可以调用的工具列表及其 JSON Schema
// (同原代码，此处省略重复注释)
func getKnowledgeGraphTools() []llms.Tool {
	return []llms.Tool{
		{
			Type: "Node operation",
			Function: &llms.FunctionDefinition{
				Name:        "search_nodes",
				Description: "当需要根据关键词在知识图谱中搜索相关节点时调用此工具",
				Parameters: map[string]any{
					"type": "object",
					"properties": map[string]any{
						"query": map[string]any{
							"type":        "string",
							"description": "搜索的关键词",
						},
					},
					"required": []string{"query"},
				},
			},
		},
		{
			Type: "Node operation",
			Function: &llms.FunctionDefinition{
				Name:        "get_neighbors",
				Description: "当需要查找某个节点的相邻节点（即与该节点有直接连线的节点）时调用此工具",
				Parameters: map[string]any{
					"type": "object",
					"properties": map[string]any{
						"node_id": map[string]any{
							"type":        "string",
							"description": "目标节点的唯一 ID",
						},
					},
					"required": []string{"node_id"},
				},
			},
		},
		{
			Type: "Node operation",
			Function: &llms.FunctionDefinition{
				Name:        "create_node",
				Description: "当需要在知识图谱中创建一个新的节点时调用此工具",
				Parameters: map[string]any{
					"type": "object",
					"properties": map[string]any{
						"id":    map[string]any{"type": "string", "description": "节点的唯一ID，格式通常为 node_xxx"},
						"label": map[string]any{"type": "string", "description": "节点的名称/标签，例如：人工智能"},
						"type":  map[string]any{"type": "string", "description": "节点的类型，例如：概念、实体"},
						"x":     map[string]any{"type": "number", "description": "节点在画布上的 X 坐标 (可选，默认0)"},
						"y":     map[string]any{"type": "number", "description": "节点在画布上的 Y 坐标 (可选，默认0)"},
						"properties": map[string]any{
							"type":        "object",
							"description": "节点的额外属性键值对 (可选)",
						},
					},
					"required": []string{"id", "label", "type"},
				},
			},
		},
		{
			Type: "Edge operation",
			Function: &llms.FunctionDefinition{
				Name:        "create_edge",
				Description: "当需要在知识图谱中创建一条连线（边）连接两个节点时调用此工具",
				Parameters: map[string]any{
					"type": "object",
					"properties": map[string]any{
						"id":     map[string]any{"type": "string", "description": "边的唯一ID，格式通常为 edge_xxx"},
						"source": map[string]any{"type": "string", "description": "起始节点的 ID"},
						"target": map[string]any{"type": "string", "description": "目标节点的 ID"},
						"label":  map[string]any{"type": "string", "description": "连线的名称/标签，例如：包含、依赖于"},
						"type":   map[string]any{"type": "string", "description": "连线的类型 (可选)，例如：CONTAINS"},
						"properties": map[string]any{
							"type":        "object",
							"description": "边的额外属性键值对 (可选)",
						},
					},
					"required": []string{"id", "source", "target", "label"},
				},
			},
		},
		{
			Type: "Node operation",
			Function: &llms.FunctionDefinition{
				Name:        "delete_node",
				Description: "当需要从知识图谱中删除一个现有节点时调用此工具",
				Parameters: map[string]any{
					"type": "object",
					"properties": map[string]any{
						"id": map[string]any{
							"type":        "string",
							"description": "要删除的节点的唯一 ID",
						},
					},
					"required": []string{"id"},
				},
			},
		},
		{
			Type: "Edge operation",
			Function: &llms.FunctionDefinition{
				Name:        "delete_edge",
				Description: "当需要从知识图谱中删除一条连线（边）时调用此工具",
				Parameters: map[string]any{
					"type": "object",
					"properties": map[string]any{
						"id": map[string]any{
							"type":        "string",
							"description": "要删除的边的唯一 ID",
						},
					},
					"required": []string{"id"},
				},
			},
		},
	}
}

// executeToolCall 执行具体的工具调用逻辑
// (同原代码，此处省略重复注释)
func executeToolCall(toolCall llms.ToolCall, graphSvc KnowledgeGraphToolService) (string, error) {
	var result any
	var err error

	switch toolCall.FunctionCall.Name {
	case "search_nodes":
		var args struct {
			Query string `json:"query"`
		}
		if err := json.Unmarshal([]byte(toolCall.FunctionCall.Arguments), &args); err != nil {
			return "", fmt.Errorf("invalid arguments for search_nodes: %w", err)
		}
		result = graphSvc.SearchNodes(args.Query)

	case "get_neighbors":
		var args struct {
			NodeID string `json:"node_id"`
		}
		if err := json.Unmarshal([]byte(toolCall.FunctionCall.Arguments), &args); err != nil {
			return "", fmt.Errorf("invalid arguments for get_neighbors: %w", err)
		}
		result, _ = graphSvc.GetNeighbors(args.NodeID)

	case "create_node":
		var args model.CreateNodeRequest
		if err := json.Unmarshal([]byte(toolCall.FunctionCall.Arguments), &args); err != nil {
			return "", fmt.Errorf("invalid arguments for create_node: %w", err)
		}
		result, err = graphSvc.CreateNode(args)
		if err != nil {
			return "", err
		}

	case "create_edge":
		var args model.CreateEdgeRequest
		if err := json.Unmarshal([]byte(toolCall.FunctionCall.Arguments), &args); err != nil {
			return "", fmt.Errorf("invalid arguments for create_edge: %w", err)
		}
		result, err = graphSvc.CreateEdge(args)
		if err != nil {
			return "", err
		}

	case "delete_node":
		var args struct {
			ID string `json:"id"`
		}
		if err := json.Unmarshal([]byte(toolCall.FunctionCall.Arguments), &args); err != nil {
			return "", fmt.Errorf("invalid arguments for delete_node: %w", err)
		}
		err = graphSvc.DeleteNode(args.ID)
		if err != nil {
			return "", err
		}
		result = map[string]string{"status": "success", "message": fmt.Sprintf("Node %s deleted successfully", args.ID)}

	case "delete_edge":
		var args struct {
			ID string `json:"id"`
		}
		if err := json.Unmarshal([]byte(toolCall.FunctionCall.Arguments), &args); err != nil {
			return "", fmt.Errorf("invalid arguments for delete_edge: %w", err)
		}
		err = graphSvc.DeleteEdge(args.ID)
		if err != nil {
			return "", err
		}
		result = map[string]string{"status": "success", "message": fmt.Sprintf("Edge %s deleted successfully", args.ID)}

	default:
		return "", fmt.Errorf("unknown tool: %s", toolCall.FunctionCall.Name)
	}

	resBytes, _ := json.Marshal(result)
	return string(resBytes), nil
}

// ChatWithTools 带有工具调用能力的 Agent 入口
// 参考了 Client.go 的上下文管理模式，加入系统提示词及 Agent Loop 循环调用
func (a *Agent) ChatWithTools(ctx context.Context, systemPrompt string, messages []llms.MessageContent) (string, error) {
	tools := getKnowledgeGraphTools()

	// 1. 构造系统提示词消息，并插入到消息列表头部
	systemMessage := llms.TextParts(llms.ChatMessageTypeSystem, systemPrompt)
	allMessages := append([]llms.MessageContent{systemMessage}, messages...)

	// 设定最大循环次数，防止死循环
	maxIterations := 10

	for i := 0; i < maxIterations; i++ {
		// 2. 调用大模型，携带上下文和可用工具
		resp, err := a.llm.GenerateContent(ctx, allMessages, llms.WithTools(tools))
		if err != nil {
			log.Printf("[ERROR] Agent GenerateContent failed: %v", err)
			return "", fmt.Errorf("agent generate content error: %w", err)
		}

		if resp == nil || len(resp.Choices) == 0 {
			log.Printf("[WARN] Agent response is empty or has no choices")
			return "", fmt.Errorf("no content choices returned")
		}

		choice := resp.Choices[0]

		// 3. 如果没有工具调用，说明大模型已经生成了最终回复，退出循环
		if len(choice.ToolCalls) == 0 {
			log.Printf("[INFO] Agent chat completed successfully, choice length: %d", len(choice.Content))
			return choice.Content, nil
		}

		// 4. 处理工具调用：需要将 AI 的工具调用指令追加到上下文中
		aiMessageParts := []llms.ContentPart{}
		if choice.Content != "" {
			aiMessageParts = append(aiMessageParts, llms.TextPart(choice.Content))
		}
		for _, tc := range choice.ToolCalls {
			aiMessageParts = append(aiMessageParts, tc)
		}
		allMessages = append(allMessages, llms.MessageContent{
			Role:  llms.ChatMessageTypeAI,
			Parts: aiMessageParts,
		})

		// 5. 依次执行每个工具调用，并将结果作为 Tool 消息追加到上下文
		for _, tc := range choice.ToolCalls {
			log.Printf("[INFO] Executing tool call: %s, Args: %s", tc.FunctionCall.Name, tc.FunctionCall.Arguments)

			toolResult, err := executeToolCall(tc, a.graphSvc)
			if err != nil {
				log.Printf("[WARN] Tool execution failed: %v", err)
				// 即使工具执行失败，也将错误信息返回给 LLM，让其进行自我修正或回复用户
				toolResult = fmt.Sprintf(`{"error": "%s"}`, err.Error())
			}

			log.Printf("[INFO] Tool call %s completed, result length: %d", tc.FunctionCall.Name, len(toolResult))

			// 构造工具调用结果并追加到上下文历史
			allMessages = append(allMessages, llms.MessageContent{
				Role: llms.ChatMessageTypeTool,
				Parts: []llms.ContentPart{
					llms.ToolCallResponse{
						ToolCallID: tc.ID,
						Name:       tc.FunctionCall.Name,
						Content:    toolResult,
					},
				},
			})
		}
	}

	return "", fmt.Errorf("agent reached maximum tool call iterations (%d)", maxIterations)
}