//
//  Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
//
//  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.
//

// QAChunker extracts question-answer pairs from parsed content.
//
// Input formats and extraction strategies:
//   - Text (txt, csv)  → delimiter-based Q&A (comma or tab)
//   - Markdown (md)    → heading-based Q&A
//   - HTML (xlsx/xls)  → table-based Q&A (first two columns)
//   - JSON (pdf, docx) → delimiter-based on structured text sections
//
// Every Q&A pair becomes a single chunk with content_with_weight
// formatted as "Question: {q}\tAnswer: {a}".
package chunker

import (
	"context"
	"encoding/csv"
	"fmt"
	"html"
	"regexp"
	"strings"

	"ragflow/internal/agent/runtime"
	"ragflow/internal/ingestion/component/schema"
	"ragflow/internal/tokenizer"
)

const ComponentNameQAChunker = "QAChunker"

type qaChunkerParam struct{}

func (p *qaChunkerParam) Update(conf map[string]any) {}

func (qaChunkerParam) Defaults() qaChunkerParam { return qaChunkerParam{} }

func (qaChunkerParam) Validate() error { return nil }

type QAChunkerComponent struct {
	name  string
	param qaChunkerParam
}

func NewQAChunker(params map[string]any) (runtime.Component, error) {
	p := qaChunkerParam{}.Defaults()
	(&p).Update(params)
	if err := p.Validate(); err != nil {
		return nil, err
	}
	return &QAChunkerComponent{
		name:  ComponentNameQAChunker,
		param: p,
	}, nil
}
func (c *QAChunkerComponent) Inputs() map[string]string { return ChunkerInputs }

func (c *QAChunkerComponent) Outputs() map[string]string { return ChunkerOutputs }

func (c *QAChunkerComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) {
	return c.invoke(ctx, inputs)
}

func (c *QAChunkerComponent) invoke(_ context.Context, inputs map[string]any) (map[string]any, error) {
	if inputs == nil {
		return emptyOutputs(), nil
	}
	upstream, err := decodeChunkerFromUpstream(inputs)
	if err != nil {
		return map[string]any{
			"output_format": "chunks",
			"chunks":        []map[string]any{},
			"_ERROR":        fmt.Sprintf("Input error: %v", err),
		}, nil
	}

	var qaPairs []qaPair
	switch upstream.OutputFormat {
	case schema.PayloadFormatHTML:
		qaPairs = extractQATable(stringPtrVal(upstream.HTMLResult))
	case schema.PayloadFormatMarkdown:
		qaPairs = extractQAMarkdown(stringPtrVal(upstream.MarkdownResult))
	case schema.PayloadFormatText:
		qaPairs = extractQAText(stringPtrVal(upstream.TextResult))
	default:
		qaPairs = extractQAJSON(upstream.JSONResult)
	}

	chunks := make([]schema.ChunkDoc, 0, len(qaPairs))
	for _, pair := range qaPairs {
		contentLTKS, _ := tokenizer.Tokenize(pair.Question)
		contentSMLTKS, _ := tokenizer.FineGrainedTokenize(contentLTKS)
		chunk := schema.ChunkDoc{
			ContentWithWeight: fmt.Sprintf("Question: %s\tAnswer: %s", rmQAPrefix(pair.Question), rmQAPrefix(pair.Answer)),
			DocType:           "text",
			ContentLtks:       contentLTKS,
			ContentSmLtks:     contentSMLTKS,
		}
		chunks = append(chunks, chunk)
	}

	return chunkOutputs(chunks), nil
}

type qaPair struct {
	Question string
	Answer   string
}

var rmQAPrefixRe = regexp.MustCompile(`^(问题|答案|回答|user|assistant|Q|A|Question|Answer|问|答)[\t:： ]+`)

func rmQAPrefix(txt string) string {
	return strings.TrimSpace(rmQAPrefixRe.ReplaceAllString(txt, ""))
}

func stringPtrVal(s *string) string {
	if s == nil {
		return ""
	}
	return *s
}

// ---------------------------------------------------------------------------
// HTML / spreadsheet QA extraction
// ---------------------------------------------------------------------------

var htmlTR = regexp.MustCompile(`(?i)<tr[^>]*>(.*?)</tr>`)
var htmlTD = regexp.MustCompile(`(?i)<t[dh][^>]*>(.*?)</t[dh]>`)
var htmlTag = regexp.MustCompile(`<[^>]+>`)

func extractQATable(htmlStr string) []qaPair {
	if htmlStr == "" {
		return nil
	}
	rows := htmlTR.FindAllStringSubmatch(htmlStr, -1)
	pairs := make([]qaPair, 0, len(rows))
	for _, row := range rows {
		cells := htmlTD.FindAllStringSubmatch(row[1], -1)
		var texts []string
		for _, cell := range cells {
			t := html.UnescapeString(htmlTag.ReplaceAllString(cell[1], ""))
			t = strings.TrimSpace(t)
			if t != "" {
				texts = append(texts, t)
			}
		}
		if len(texts) >= 2 {
			pairs = append(pairs, qaPair{Question: texts[0], Answer: texts[1]})
		}
	}
	return pairs
}

// ---------------------------------------------------------------------------
// Markdown QA extraction
// ---------------------------------------------------------------------------

var mdHeading = regexp.MustCompile(`^(#{1,6})\s+`)

func extractQAMarkdown(md string) []qaPair {
	if md == "" {
		return nil
	}
	lines := strings.Split(md, "\n")
	var pairs []qaPair
	var questionStack, levelStack []string
	var answer []string
	codeBlock := false

	flushAnswer := func() {
		joined := strings.TrimSpace(strings.Join(answer, "\n"))
		if joined != "" && len(questionStack) > 0 {
			sumQ := strings.Join(questionStack, "\n")
			pairs = append(pairs, qaPair{Question: sumQ, Answer: joined})
		}
		answer = nil
	}

	for _, line := range lines {
		trimmed := strings.TrimSpace(line)
		if strings.HasPrefix(trimmed, "```") {
			codeBlock = !codeBlock
		}
		if codeBlock {
			answer = append(answer, line)
			continue
		}

		m := mdHeading.FindStringSubmatch(line)
		if m == nil || len(m[1]) > 6 {
			answer = append(answer, line)
			continue
		}

		flushAnswer()
		level := m[1]
		question := strings.TrimSpace(line[len(m[0]):])

		for len(levelStack) > 0 && len(level) <= len(levelStack[len(levelStack)-1]) {
			questionStack = questionStack[:len(questionStack)-1]
			levelStack = levelStack[:len(levelStack)-1]
		}
		questionStack = append(questionStack, question)
		levelStack = append(levelStack, level)
	}
	flushAnswer()
	return pairs
}

// ---------------------------------------------------------------------------
// Text / delimiter-based QA extraction (txt, csv)
// ---------------------------------------------------------------------------

func extractQAText(text string) []qaPair {
	if text == "" {
		return nil
	}
	lines := strings.Split(text, "\n")

	delimiter := detectDelimiter(lines)

	var pairs []qaPair
	var question, answer string

	for _, line := range lines {
		if strings.TrimSpace(line) == "" {
			continue
		}
		parts := splitQA(line, delimiter)
		if len(parts) != 2 {
			if question != "" {
				answer += "\n" + line
			}
			continue
		}
		if question != "" && answer != "" {
			pairs = append(pairs, qaPair{Question: strings.TrimSpace(question), Answer: strings.TrimSpace(answer)})
		}
		question = parts[0]
		answer = parts[1]
	}
	if question != "" {
		pairs = append(pairs, qaPair{Question: strings.TrimSpace(question), Answer: strings.TrimSpace(answer)})
	}
	return pairs
}

func detectDelimiter(lines []string) string {
	comma, tab := 0, 0
	for _, line := range lines {
		if len(strings.Split(line, ",")) == 2 {
			comma++
		}
		if len(strings.Split(line, "\t")) == 2 {
			tab++
		}
	}
	if tab >= comma {
		return "\t"
	}
	return ","
}

func splitQA(line, delimiter string) []string {
	if delimiter == "\t" {
		parts := strings.Split(line, "\t")
		if len(parts) == 2 {
			return parts
		}
		return []string{line}
	}
	r := csv.NewReader(strings.NewReader(line))
	r.Comma = ','
	r.LazyQuotes = true
	records, err := r.Read()
	if err != nil || len(records) != 2 {
		return []string{line}
	}
	return records
}

// ---------------------------------------------------------------------------
// JSON / structured QA extraction
// ---------------------------------------------------------------------------

func extractQAJSON(items []schema.ChunkDoc) []qaPair {
	var pairs []qaPair
	for _, item := range items {
		txt, _ := itemText(item)
		if txt == "" {
			continue
		}
		tmp := extractQAText(txt)
		pairs = append(pairs, tmp...)
	}
	return pairs
}

func init() {
	MustRegisterChunker(ComponentNameQAChunker)
}