package service

import (
	"context"
	"crypto/sha256"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"testing"

	chatpipeline "github.com/Tencent/WeKnora/internal/application/service/chat_pipeline"
	"github.com/Tencent/WeKnora/internal/models/chat"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/Tencent/WeKnora/internal/types/interfaces"
	"github.com/hibiken/asynq"
)

func TestNativeGraphReviewedSourceStopsBeforeModelOnChange(t *testing.T) {
	const source = "Trust supports learning."
	for _, changed := range []string{"unchanged", "content", "knowledge", "knowledge base", "chunk", "tenant", "empty identity", "uppercase hash", "disabled", "deleted", "kind", "document deleted", "document processing", "document tenant", "document base", "document identity", "document unavailable", "document repository unavailable", "base deleted", "base tenant", "base identity"} {
		t.Run(changed, func(t *testing.T) {
			chunk := &types.Chunk{ID: "chunk", TenantID: 7, KnowledgeBaseID: "kb", KnowledgeID: "doc", Content: source, IsEnabled: true, ChunkType: types.ChunkTypeText}
			expected := &types.GraphSourceIdentity{KnowledgeBaseID: "kb", KnowledgeID: "doc", ContentSHA256: fmt.Sprintf("%x", sha256.Sum256([]byte(source)))}
			document := &types.Knowledge{ID: "doc", TenantID: 7, KnowledgeBaseID: "kb", ParseStatus: types.ParseStatusCompleted}
			knowledge := &extractionKnowledge{document: document}
			base := &types.KnowledgeBase{ID: "kb", TenantID: 7, ExtractConfig: &types.ExtractConfig{Enabled: true}}
			switch changed {
			case "content":
				chunk.Content = "New source."
			case "knowledge":
				chunk.KnowledgeID = "other"
			case "knowledge base":
				chunk.KnowledgeBaseID = "other"
			case "chunk":
				chunk.ID = "other"
			case "tenant":
				chunk.TenantID = 8
			case "empty identity":
				expected.KnowledgeID = ""
			case "uppercase hash":
				expected.ContentSHA256 = strings.ToUpper(expected.ContentSHA256)
			case "disabled":
				chunk.IsEnabled = false
			case "deleted":
				chunk.DeletedAt.Valid = true
			case "kind":
				chunk.ChunkType = types.ChunkTypeImageOCR
			case "document deleted":
				document.DeletedAt.Valid = true
			case "document processing":
				document.ParseStatus = types.ParseStatusProcessing
			case "document tenant":
				document.TenantID = 8
			case "document base":
				document.KnowledgeBaseID = "other"
			case "document identity":
				document.ID = "other"
			case "document unavailable":
				knowledge.lookupErr = errors.New("unavailable")
			case "base deleted":
				base.DeletedAt.Valid = true
			case "base tenant":
				base.TenantID = 8
			case "base identity":
				base.ID = "other"
			}
			model := &extractionChat{response: `[{"entity1":"Trust","entity2":"learning","relation":"supports","evidence":"Trust supports learning."}]`}
			writer := &extractionWriter{}
			service := &ChunkExtractService{template: &types.PromptTemplateStructured{Description: "Extract graph"}, modelService: extractionModels{model: model}, knowledgeBaseRepo: extractionKB{base: base}, knowledgeRepo: knowledge, chunkRepo: extractionChunks{chunk: chunk}, graphEngine: writer}
			if changed == "document repository unavailable" {
				service.knowledgeRepo = nil
			}
			service.executionDB = graphBudgetDB(t)
			payload, _ := json.Marshal(types.ExtractChunkPayload{TenantID: 7, ChunkID: "chunk", ModelID: "model", ExpectedSource: expected})
			err := service.handle(context.Background(), asynq.NewTask(types.TypeChunkExtract, payload), "task")
			if changed == "unchanged" {
				if err != nil || model.calls != 1 || !writer.complete {
					t.Fatalf("reviewed source rejected: %v", err)
				}
			} else if !errors.Is(err, asynq.SkipRetry) || model.calls != 0 || writer.saved != nil {
				t.Fatalf("changed source reached model or persistence: %v calls=%d", err, model.calls)
			}
		})
	}
}

func TestNativeGraphExtractionDoesNotRequireNeo4j(t *testing.T) {
	t.Setenv("NEO4J_ENABLE", "false")
	t.Setenv("LEGACY_GRAPH_ENABLE", "true")
	if !graphExtractionEnabled() {
		t.Fatal("native extraction was silently skipped")
	}
	t.Setenv("LEGACY_GRAPH_ENABLE", "false")
	if graphExtractionEnabled() {
		t.Fatal("unconfigured graph extraction enabled")
	}
	t.Setenv("NEO4J_ENABLE", "true")
	if !graphExtractionEnabled() {
		t.Fatal("Neo4j compatibility lost")
	}
}

type extractionChatBase interface{ chat.Chat }
type extractionChat struct {
	extractionChatBase
	calls    int
	response string
}

func (m *extractionChat) Chat(ctx context.Context, messages []chat.Message, opts *chat.ChatOptions) (*types.ChatResponse, error) {
	m.calls++
	return &types.ChatResponse{Content: m.response}, nil
}

type extractionModels struct {
	interfaces.ModelService
	model chat.Chat
}

func (m extractionModels) GetChatModel(context.Context, string) (chat.Chat, error) {
	return m.model, nil
}

type extractionKB struct {
	interfaces.KnowledgeBaseRepository
	base *types.KnowledgeBase
}

func (r extractionKB) GetKnowledgeBaseByID(context.Context, string) (*types.KnowledgeBase, error) {
	if r.base != nil {
		return r.base, nil
	}
	return &types.KnowledgeBase{ID: "kb", TenantID: 7, ExtractConfig: &types.ExtractConfig{Enabled: true}}, nil
}

type extractionChunks struct {
	interfaces.ChunkRepository
	chunk *types.Chunk
}

type extractionKnowledge struct {
	interfaces.KnowledgeRepository
	drained   int
	document  *types.Knowledge
	lookupErr error
}

func (r *extractionKnowledge) GetKnowledgeByIDOnly(context.Context, string) (*types.Knowledge, error) {
	if r.lookupErr != nil {
		return nil, r.lookupErr
	}
	if r.document != nil {
		return r.document, nil
	}
	return &types.Knowledge{ID: "doc", TenantID: 7, KnowledgeBaseID: "kb", ParseStatus: types.ParseStatusCompleted}, nil
}

func (r *extractionKnowledge) FinalizeSubtask(context.Context, string) (int, bool, error) {
	r.drained++
	return 0, true, nil
}

func (r extractionChunks) GetChunkByID(context.Context, uint64, string) (*types.Chunk, error) {
	copy := *r.chunk
	return &copy, nil
}

type extractionWriter struct {
	interfaces.RetrieveGraphRepository
	saved    *types.GraphData
	source   *types.Chunk
	complete bool
	fail     bool
}

func (w *extractionWriter) HasChunkGraph(context.Context, *types.Chunk) (bool, error) {
	return w.complete, nil
}
func (w *extractionWriter) AddChunkGraph(ctx context.Context, ns types.NameSpace, c *types.Chunk, g *types.GraphData) error {
	if w.fail {
		return fmt.Errorf("persistence failure")
	}
	if ns.Knowledge != "doc" || types.MustTenantIDFromContext(ctx) != 7 {
		return fmt.Errorf("missing ownership scope")
	}
	w.saved, w.source, w.complete = g, c, true
	return nil
}

func TestNativeGraphTaskExtractsBeforeCheckpointing(t *testing.T) {
	for _, tc := range []struct {
		name, response string
		fail, wantErr  bool
	}{
		{"actual entities and relationship", `[{"entity":"Trust","entity_attributes":["concept"]},{"entity":"Learning"},{"entity1":"Trust","entity2":"Learning","relation":"supports","evidence":"Trust supports learning."}]`, false, false},
		{"missing relation evidence", `[{"entity1":"Trust","entity2":"Learning","relation":"supports"}]`, false, true},
		{"fabricated relation evidence", `[{"entity1":"Trust","entity2":"Learning","relation":"supports","evidence":"Invented source."}]`, false, true},
		{"explicit empty needs review", `[]`, false, true},
		{"entities without relations need review", `[{"entity":"Trust"}]`, false, true},
		{"unrecognized response", `{"error":"provider rejected extraction"}`, false, true},
		{"database error", `[{"entity1":"Trust","entity2":"Learning","relation":"supports","evidence":"Trust supports learning."}]`, true, true},
	} {
		t.Run(tc.name, func(t *testing.T) {
			model := &extractionChat{response: tc.response}
			writer := &extractionWriter{fail: tc.fail}
			service := &ChunkExtractService{template: &types.PromptTemplateStructured{Description: "Extract grounded entities and relationships."}, modelService: extractionModels{model: model}, knowledgeBaseRepo: extractionKB{}, chunkRepo: extractionChunks{chunk: &types.Chunk{ID: "chunk", TenantID: 7, KnowledgeBaseID: "kb", KnowledgeID: "doc", Content: "Trust supports learning."}}, graphEngine: writer}
			payload, _ := json.Marshal(types.ExtractChunkPayload{TenantID: 7, ChunkID: "chunk", ModelID: "model"})
			task := asynq.NewTask(types.TypeChunkExtract, payload)
			err := service.Handle(context.Background(), task)
			if (err != nil) != tc.wantErr {
				t.Fatalf("unexpected task result: %v", err)
			}
			if tc.wantErr {
				if writer.complete {
					t.Fatal("failed task wrote a completion checkpoint")
				}
				if tc.name == "explicit empty needs review" || tc.name == "entities without relations need review" {
					if !errors.Is(err, chatpipeline.ErrGraphNeedsReview) || !errors.Is(err, asynq.SkipRetry) {
						t.Fatalf("review result should stop automatic retries: %v", err)
					}
				}
				return
			}
			if writer.saved == nil || writer.source == nil || writer.source.Content != "Trust supports learning." {
				t.Fatal("model output/source not passed to durable writer")
			}
			if len(writer.saved.Node) > 0 && writer.saved.Node[0].Chunks[0] != "chunk" {
				t.Fatal("missing evidence chunk")
			}
			if len(writer.saved.Relation) > 0 && writer.saved.Relation[0].Evidence != writer.source.Content {
				t.Fatal("relation evidence was not passed to durable writer")
			}
			if err := service.Handle(context.Background(), task); err != nil {
				t.Fatal(err)
			}
			if model.calls != 1 {
				t.Fatal("completed retry billed the model again")
			}
		})
	}
}

func TestNativeGraphTaskDoesNotCheckpointContradictedProductClaim(t *testing.T) {
	const source = "私蜜妍+蜜桃初醒套盒系列\n过敏者慎甩；粘膜破损处禁用；不得用于性生活中对性病的预防。"
	model := &extractionChat{response: `[{"entity1":"私蜜妍+蜜桃初醒套盒系列","entity2":"粘膜破损处","relation":"适用于","evidence":"粘膜破损处禁用"}]`}
	writer := &extractionWriter{}
	service := &ChunkExtractService{
		template:     &types.PromptTemplateStructured{Description: "Extract grounded entities and relationships."},
		modelService: extractionModels{model: model}, knowledgeBaseRepo: extractionKB{},
		chunkRepo:   extractionChunks{chunk: &types.Chunk{ID: "chunk", TenantID: 7, KnowledgeBaseID: "kb", KnowledgeID: "doc", Content: source}},
		graphEngine: writer,
	}
	payload, _ := json.Marshal(types.ExtractChunkPayload{TenantID: 7, ChunkID: "chunk", ModelID: "model"})
	if err := service.Handle(context.Background(), asynq.NewTask(types.TypeChunkExtract, payload)); err == nil {
		t.Fatal("contradicted claim was accepted")
	}
	if writer.saved != nil || writer.complete {
		t.Fatal("contradicted claim wrote a graph or completion checkpoint")
	}
}

func TestNativeGraphReviewDrainsParentWithoutRetryOrCheckpoint(t *testing.T) {
	model := &extractionChat{response: `[]`}
	writer := &extractionWriter{}
	knowledge := &extractionKnowledge{}
	service := &ChunkExtractService{
		template:     &types.PromptTemplateStructured{Description: "Extract grounded relations."},
		modelService: extractionModels{model: model}, knowledgeBaseRepo: extractionKB{},
		knowledgeRepo: knowledge,
		chunkRepo:     extractionChunks{chunk: &types.Chunk{ID: "chunk", TenantID: 7, KnowledgeBaseID: "kb", KnowledgeID: "doc", Content: "Trust supports learning."}},
		graphEngine:   writer,
	}
	payload, _ := json.Marshal(types.ExtractChunkPayload{TenantID: 7, KnowledgeID: "doc", ChunkID: "chunk", ModelID: "model"})
	err := service.Handle(context.Background(), asynq.NewTask(types.TypeChunkExtract, payload))
	if !errors.Is(err, asynq.SkipRetry) || !errors.Is(err, chatpipeline.ErrGraphNeedsReview) {
		t.Fatalf("unexpected review outcome: %v", err)
	}
	if knowledge.drained != 1 || writer.complete || writer.saved != nil {
		t.Fatalf("review left the parent pending or wrote a checkpoint: drained=%d complete=%t", knowledge.drained, writer.complete)
	}
	if model.calls != 2 {
		t.Fatalf("expected one bounded correction, got %d calls", model.calls)
	}
}