/*
 * Copyright 2026 CloudWeGo Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

// Package rag provides an answer_from_document tool backed by a compose.Workflow.
//
// The workflow uses field mapping to share the user's question across non-adjacent
// nodes (score, answer) without threading it through every intermediate output type:
//
//	START{FilePath, Question}
//	  │ (data via WithNoDirectDependency)──────────────────────────────────────────┐
//	  ▼                                                                            │ Question
//	[load]  os.ReadFile → []*schema.Document                                       │
//	  ▼                                                                            │
//	[chunk] paragraph splitter → []*schema.Document                               │
//	  ▼  Chunks ─────────────────────────────────────────────────────────────► [score]
//	                                                                               │ []scoredChunk
//	                                                                               ▼
//	                                                                           [filter]  top-k
//	                                                                               │ TopK (may be empty)
//	                                                                               ▼
//	                                                                           [answer] ◄─ Question (START)
//	                                                                    (synthesize or not_found inline)
//	                                                                               │
//	                                                                              END
//
// The [score] node wraps a BatchNode whose inner workflow scores each chunk with
// a ChatModel call in parallel (MaxConcurrency=5).
package rag

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"sort"
	"strings"

	"github.com/cloudwego/eino/components/model"
	"github.com/cloudwego/eino/components/tool"
	"github.com/cloudwego/eino/compose"
	"github.com/cloudwego/eino/schema"

	"github.com/cloudwego/eino-examples/adk/common/tool/graphtool"
	"github.com/cloudwego/eino-examples/compose/batch/batch"
)

// Input is the tool call argument struct. Its JSON tags are used by utils.InferTool
// to generate the tool's parameter schema automatically.
type Input struct {
	FilePath string `json:"file_path" jsonschema:"description=Absolute path to the uploaded document file"`
	Question string `json:"question"  jsonschema:"description=The question to answer from the document"`
}

// Output is the structured result returned by the tool.
type Output struct {
	Answer  string   `json:"answer"`
	Sources []string `json:"sources"` // key excerpts used to produce the answer
}

// scoreTask is the per-chunk input fed into the inner BatchNode workflow.
type scoreTask struct {
	Text     string
	Question string
}

// scoredChunk is the per-chunk result produced by the inner BatchNode workflow.
type scoredChunk struct {
	Text    string
	Score   int    // 0–10 relevance to the question
	Excerpt string // most relevant sentence or phrase from this chunk
}

// scoreIn is the input to the outer "score" Lambda node.
// It is assembled by field mapping from two sources:
//   - Chunks: full output of "chunk" node ([]*schema.Document)
//   - Question: Question field of START (Input)
type scoreIn struct {
	Chunks   []*schema.Document
	Question string
}

// synthIn is the input to the "synthesize" Lambda node.
// It is assembled by field mapping from two sources:
//   - TopK: full output of "filter" node ([]scoredChunk)
//   - Question: Question field of START (Input)
type synthIn struct {
	TopK     []scoredChunk
	Question string
}

// BuildTool constructs the answer_from_document tool backed by the RAG workflow.
// It uses graphtool.NewInvokableGraphTool, which compiles the workflow per invocation
// and supports interrupt/resume via a built-in checkpoint store.
func BuildTool(ctx context.Context, cm model.BaseChatModel) (tool.BaseTool, error) {
	wf := buildWorkflow(cm)
	return graphtool.NewInvokableGraphTool[Input, Output](
		wf,
		"answer_from_document",
		"Search a large uploaded document for content relevant to a question and synthesize a "+
			"cited answer from the most relevant passages. "+
			"Use this instead of read_file when the document may be too large to fit in context.",
	)
}

// buildWorkflow constructs the RAG compose.Workflow (uncompiled).
// graphtool.NewInvokableGraphTool compiles it per invocation.
func buildWorkflow(cm model.BaseChatModel) *compose.Workflow[Input, Output] {
	scoreWF := newScoreWorkflow(cm)
	scorer := batch.NewBatchNode(&batch.NodeConfig[scoreTask, scoredChunk]{
		Name:           "ChunkScorer",
		InnerTask:      scoreWF,
		MaxConcurrency: 5,
	})

	wf := compose.NewWorkflow[Input, Output]()

	// load: read file from disk, emit a single Document.
	wf.AddLambdaNode("load", compose.InvokableLambda(
		func(ctx context.Context, in Input) ([]*schema.Document, error) {
			data, err := os.ReadFile(in.FilePath)
			if err != nil {
				return nil, fmt.Errorf("read %q: %w", in.FilePath, err)
			}
			return []*schema.Document{{Content: string(data)}}, nil
		},
	)).AddInput(compose.START)

	// chunk: split each Document into ~800-char pieces.
	wf.AddLambdaNode("chunk", compose.InvokableLambda(
		func(ctx context.Context, docs []*schema.Document) ([]*schema.Document, error) {
			var out []*schema.Document
			for _, d := range docs {
				out = append(out, splitIntoChunks(d.Content, 800)...)
			}
			return out, nil
		},
	)).AddInput("load")

	// score: score each chunk against the question in parallel via BatchNode.
	// Chunks comes from "chunk"; Question comes directly from START.
	// Both use WithNoDirectDependency because the execution order is already
	// established by the direct edges START→load→chunk→score.
	wf.AddLambdaNode("score", compose.InvokableLambda(
		func(ctx context.Context, in scoreIn) ([]scoredChunk, error) {
			tasks := make([]scoreTask, len(in.Chunks))
			for i, c := range in.Chunks {
				tasks[i] = scoreTask{Text: c.Content, Question: in.Question}
			}
			return scorer.Invoke(ctx, tasks)
		},
	)).
		AddInputWithOptions("chunk",
			[]*compose.FieldMapping{compose.ToField("Chunks")},
			compose.WithNoDirectDependency()).
		AddInputWithOptions(compose.START,
			[]*compose.FieldMapping{compose.MapFields("Question", "Question")},
			compose.WithNoDirectDependency())

	// filter: sort descending by score, keep up to top-3 chunks with score ≥ 3.
	wf.AddLambdaNode("filter", compose.InvokableLambda(
		func(ctx context.Context, scored []scoredChunk) ([]scoredChunk, error) {
			sort.Slice(scored, func(i, j int) bool {
				return scored[i].Score > scored[j].Score
			})
			const maxK = 3
			var top []scoredChunk
			for _, c := range scored {
				if c.Score < 3 {
					break
				}
				top = append(top, c)
				if len(top) == maxK {
					break
				}
			}
			return top, nil
		},
	)).AddInput("score")

	// answer: synthesize a response from top-k chunks, or return a not-found message if empty.
	// TopK comes from "filter"; Question comes directly from START.
	// Both use WithNoDirectDependency: "filter" governs execution order via its direct edge.
	wf.AddLambdaNode("answer", compose.InvokableLambda(
		func(ctx context.Context, in synthIn) (Output, error) {
			if len(in.TopK) == 0 {
				return Output{
					Answer: fmt.Sprintf("No relevant content found in the document for: %q", in.Question),
				}, nil
			}
			return synthesize(ctx, cm, in)
		},
	)).
		AddInputWithOptions("filter",
			[]*compose.FieldMapping{compose.ToField("TopK")},
			compose.WithNoDirectDependency()).
		AddInputWithOptions(compose.START,
			[]*compose.FieldMapping{compose.MapFields("Question", "Question")},
			compose.WithNoDirectDependency())

	// END receives output from answer.
	wf.End().
		AddInput("answer")

	return wf
}

// newScoreWorkflow builds the single-node inner workflow used by each BatchNode task.
// It is intentionally trivial: BatchNode provides the parallelism, not the inner graph.
func newScoreWorkflow(cm model.BaseChatModel) *compose.Workflow[scoreTask, scoredChunk] {
	wf := compose.NewWorkflow[scoreTask, scoredChunk]()
	wf.AddLambdaNode("score_chunk", compose.InvokableLambda(
		func(ctx context.Context, t scoreTask) (scoredChunk, error) {
			return scoreOneChunk(ctx, cm, t)
		},
	)).AddInput(compose.START)
	wf.End().AddInput("score_chunk")
	return wf
}

// scoreOneChunk asks the model to rate the relevance of a single chunk (0–10)
// and extract the most relevant excerpt. Parse errors are treated as score 0
// so a bad JSON response never aborts the pipeline.
func scoreOneChunk(ctx context.Context, cm model.BaseChatModel, t scoreTask) (scoredChunk, error) {
	prompt := fmt.Sprintf(`Rate how relevant the following text chunk is to the question.

Question: %s

Chunk:
%s

Reply with JSON only — no explanation, no markdown fences:
{"score": <0-10>, "excerpt": "<most relevant sentence or phrase, empty string if score is 0>"}

Score guide: 0=completely irrelevant, 3=tangentially related, 7=clearly relevant, 10=directly answers the question.`,
		t.Question, t.Text)

	resp, err := cm.Generate(ctx, []*schema.Message{schema.UserMessage(prompt)})
	if err != nil {
		// treat model error as irrelevant rather than aborting the batch
		return scoredChunk{Text: t.Text, Score: 0}, nil
	}

	content := strings.TrimSpace(resp.Content)
	// strip optional markdown code block wrapper
	content = strings.TrimPrefix(content, "```json")
	content = strings.TrimPrefix(content, "```")
	content = strings.TrimSuffix(content, "```")
	content = strings.TrimSpace(content)

	var sr struct {
		Score   int    `json:"score"`
		Excerpt string `json:"excerpt"`
	}
	if err := json.Unmarshal([]byte(content), &sr); err != nil {
		return scoredChunk{Text: t.Text, Score: 0}, nil
	}
	return scoredChunk{Text: t.Text, Score: sr.Score, Excerpt: sr.Excerpt}, nil
}

// synthesize builds a prompt from the top-k chunks and generates a cited answer.
func synthesize(ctx context.Context, cm model.BaseChatModel, in synthIn) (Output, error) {
	var sb strings.Builder
	sb.WriteString("Answer the following question using only the provided document excerpts.\n\n")
	sb.WriteString("Question: ")
	sb.WriteString(in.Question)
	sb.WriteString("\n\nDocument excerpts:\n")

	sources := make([]string, len(in.TopK))
	for i, c := range in.TopK {
		excerpt := c.Excerpt
		if excerpt == "" {
			excerpt = c.Text
		}
		sources[i] = excerpt
		fmt.Fprintf(&sb, "[%d] %s\n\n", i+1, excerpt)
	}
	sb.WriteString("Provide a clear, concise answer. Cite excerpt numbers like [1] when referencing sources.")

	resp, err := cm.Generate(ctx, []*schema.Message{schema.UserMessage(sb.String())})
	if err != nil {
		return Output{}, fmt.Errorf("synthesize: %w", err)
	}
	return Output{Answer: resp.Content, Sources: sources}, nil
}

// splitIntoChunks splits text into chunks of at most chunkSize characters,
// breaking on paragraph boundaries (\n\n) where possible, then on newlines.
func splitIntoChunks(text string, chunkSize int) []*schema.Document {
	var chunks []*schema.Document
	var buf strings.Builder

	flush := func() {
		s := strings.TrimSpace(buf.String())
		if s != "" {
			chunks = append(chunks, &schema.Document{Content: s})
		}
		buf.Reset()
	}

	for _, para := range strings.Split(text, "\n\n") {
		para = strings.TrimSpace(para)
		if para == "" {
			continue
		}
		if buf.Len()+len(para)+2 > chunkSize && buf.Len() > 0 {
			flush()
		}
		// paragraph itself exceeds chunkSize: split by line
		if len(para) > chunkSize {
			for _, line := range strings.Split(para, "\n") {
				line = strings.TrimSpace(line)
				if line == "" {
					continue
				}
				if buf.Len()+len(line)+1 > chunkSize && buf.Len() > 0 {
					flush()
				}
				if buf.Len() > 0 {
					buf.WriteByte('\n')
				}
				buf.WriteString(line)
			}
		} else {
			if buf.Len() > 0 {
				buf.WriteString("\n\n")
			}
			buf.WriteString(para)
		}
	}
	flush()
	return chunks
}