package service

import (
	"context"
	"fmt"
	"strings"

	"github.com/gochat/gochat/internal/llm"
	"github.com/gochat/gochat/internal/model"
	"github.com/pgvector/pgvector-go"
	applogger "github.com/gochat/gochat/pkg/logger"
)

// RAGService implements Retrieval-Augmented Generation for Captain knowledge base Q&A.
// Reference: Chatwoot Captain::Llm::GraphService + Captain::Tools::CaptainKnowledgeBaseTool
//
// RAG flow:
// 1. Take user question → generate embedding via LLM provider
// 2. Search assistant response embeddings in pgvector for top-k relevant FAQs
// 3. Build prompt with retrieved context + question → call LLM for answer
// 4. Return answer with source document references

// RAGQueryRequest holds the input for a RAG Q&A query.
type RAGQueryRequest struct {
	AssistantID uint   `json:"assistant_id" validate:"required"`
	Question    string `json:"question" validate:"required,min=1"`
	TopK        int    `json:"top_k,omitempty"` // number of FAQ responses to retrieve, default 5
}

// RAGQueryResult holds the answer from RAG Q&A.
type RAGQueryResult struct {
	Answer       string      `json:"answer"`
	Sources      []RAGSource `json:"sources,omitempty"`
	Confidence   float64     `json:"confidence"`
	AssistantID  uint        `json:"assistant_id"`
	Conversation uint        `json:"conversation_id,omitempty"`
}

// RAGSource references a source FAQ response used in the answer.
type RAGSource struct {
	ResponseID  uint   `json:"response_id"`
	Question    string `json:"question"`
	Answer      string `json:"answer"`
	DocumentID  uint   `json:"document_id,omitempty"`
	Score       float64 `json:"score"`
}

// AssistantRepoIface defines the repository interface RAGService needs from CaptainAssistantRepo.
type AssistantRepoIface interface {
	GetByID(ctx context.Context, id uint) (*model.CaptainAssistant, error)
}

// ResponseRepoIface defines the repository interface RAGService needs from CaptainAssistantResponseRepo.
type ResponseRepoIface interface {
	SearchByEmbedding(ctx context.Context, assistantID uint, embedding pgvector.Vector, limit int) ([]model.CaptainAssistantResponse, error)
	GetByID(ctx context.Context, id uint) (*model.CaptainAssistantResponse, error)
	Update(ctx context.Context, resp *model.CaptainAssistantResponse) error
}

// RAGService orchestrates embedding search + LLM generation for knowledge base Q&A.
type RAGService struct {
	responseRepo  ResponseRepoIface
	assistantRepo AssistantRepoIface
	llmProvider   llm.Provider
	promptBuilder *SystemPromptBuilder
}

// NewRAGService creates a new RAGService.
func NewRAGService(
	responseRepo ResponseRepoIface,
	assistantRepo AssistantRepoIface,
	llmProvider llm.Provider,
) *RAGService {
	return &RAGService{
		responseRepo:  responseRepo,
		assistantRepo: assistantRepo,
		llmProvider:   llmProvider,
		promptBuilder: NewSystemPromptBuilder(),
	}
}

// Query performs a RAG Q&A: embed question → search FAQs → generate answer.
func (s *RAGService) Query(ctx context.Context, accountID uint, req *RAGQueryRequest) (*RAGQueryResult, error) {
	if req.TopK <= 0 {
		req.TopK = 5
	}

	// Step 1: Verify assistant exists and is active
	assistant, err := s.assistantRepo.GetByID(ctx, req.AssistantID)
	if err != nil {
		return nil, fmt.Errorf("assistant not found: %w", err)
	}
	if assistant.Status != model.AssistantStatusActive {
		return nil, fmt.Errorf("assistant is not active (status: %s)", assistant.Status)
	}

	// Get assistant config for prompt building and model selection
	cfg, err := assistant.GetConfig()
	if err != nil {
		applogger.L().Warnf("RAG get assistant config: %v", err)
		cfg = &model.AssistantConfig{}
	}

	// Step 2: Generate embedding for the question
	embedResp, err := s.llmProvider.CreateEmbedding(ctx, llm.EmbeddingRequest{
		Model: "text-embedding-3-small",
		Input: []string{req.Question},
	})
	if err != nil {
		applogger.L().Errorf("RAG embedding generation: %v", err)
		return nil, fmt.Errorf("embedding generation failed: %w", err)
	}
	if len(embedResp.Data) == 0 {
		return nil, fmt.Errorf("no embedding returned for question")
	}
	questionEmbedding := embedResp.Data[0].Embedding

	// Convert float64 embedding to float32 for pgvector
	float32Emb := make([]float32, len(questionEmbedding))
	for i, v := range questionEmbedding {
		float32Emb[i] = float32(v)
	}
	pgvectorEmb := pgvector.NewVector(float32Emb)

	// Step 3: Search assistant responses by embedding similarity (pgvector cosine distance)
	results, err := s.responseRepo.SearchByEmbedding(ctx, req.AssistantID, pgvectorEmb, req.TopK)
	if err != nil {
		applogger.L().Errorf("RAG response search: %v", err)
		return nil, fmt.Errorf("FAQ search failed: %w", err)
	}

	if len(results) == 0 {
		// No relevant FAQs found — fallback to LLM without context
		return s.queryWithoutContext(ctx, assistant, cfg, req.Question)
	}

	// Step 4: Build RAG context from retrieved FAQ responses
	var contextParts []string
	var sources []RAGSource
	for i, r := range results {
		contextParts = append(contextParts, fmt.Sprintf("[FAQ %d]\nQ: %s\nA: %s", i+1, r.Question, r.Answer))

		// Estimate score (pgvector cosine distance is already sorted by closest)
		score := 1.0 - float64(i)*0.1 // rough heuristic since we don't have raw distance
		if i == 0 {
			score = 0.9 // top result is likely high-confidence
		}
		sources = append(sources, RAGSource{
			ResponseID: r.ID,
			Question:   r.Question,
			Answer:     r.Answer,
			Score:      score,
		})
		if r.DocumentableID != nil {
			sources[len(sources)-1].DocumentID = *r.DocumentableID
		}
	}
	ragContext := strings.Join(contextParts, "\n\n")

	// Step 5: Build prompt with RAG context + system prompt
	systemPrompt := s.promptBuilder.BuildAssistantPrompt(assistant, cfg)
	systemPrompt += fmt.Sprintf("\n\nKnowledge Base Context:\n%s", ragContext)
	systemPrompt += "\n\nImportant: When answering, reference the FAQ entries by number. If the FAQ context doesn't contain enough information to answer, say so honestly."

	messages := []llm.ChatMessage{
		{Role: "system", Content: systemPrompt},
		{Role: "user", Content: req.Question},
	}

	// Step 6: Call LLM for answer generation
	modelName := cfg.Model
	if modelName == "" {
		modelName = "gpt-4"
	}
	temperature := cfg.Temperature
	if temperature == 0 {
		temperature = 0.3 // lower temp for factual answers
	}

	llmResp, err := s.llmProvider.ChatCompletion(ctx, llm.ChatRequest{
		Model:       modelName,
		Messages:    messages,
		Temperature: temperature,
		MaxTokens:   1024,
	})
	if err != nil {
		applogger.L().Errorf("RAG LLM generation: %v", err)
		return nil, fmt.Errorf("answer generation failed: %w", err)
	}

	answer := ""
	confidence := 0.0
	if len(llmResp.Choices) > 0 {
		answer = llmResp.Choices[0].Message.Content
		if len(sources) > 0 {
			confidence = sources[0].Score
		}
	}

	return &RAGQueryResult{
		Answer:      answer,
		Sources:     sources,
		Confidence:  confidence,
		AssistantID: req.AssistantID,
	}, nil
}

// queryWithoutContext handles the case when no relevant FAQ responses are found.
// Falls back to generating an answer purely from the assistant's configuration.
func (s *RAGService) queryWithoutContext(ctx context.Context, assistant *model.CaptainAssistant, cfg *model.AssistantConfig, question string) (*RAGQueryResult, error) {
	systemPrompt := s.promptBuilder.BuildAssistantPrompt(assistant, cfg)
	systemPrompt += "\n\nNote: No relevant FAQ entries were found for this question. Answer based on your general knowledge, but indicate that the answer may not be specific to the product."

	modelName := cfg.Model
	if modelName == "" {
		modelName = "gpt-4"
	}
	temperature := cfg.Temperature
	if temperature == 0 {
		temperature = 0.5
	}

	llmResp, err := s.llmProvider.ChatCompletion(ctx, llm.ChatRequest{
		Model:       modelName,
		Messages: []llm.ChatMessage{
			{Role: "system", Content: systemPrompt},
			{Role: "user", Content: question},
		},
		Temperature: temperature,
		MaxTokens:   512,
	})
	if err != nil {
		return nil, fmt.Errorf("fallback LLM generation failed: %w", err)
	}

	answer := ""
	if len(llmResp.Choices) > 0 {
		answer = llmResp.Choices[0].Message.Content
	}

	return &RAGQueryResult{
		Answer:      answer,
		Confidence:  0.3, // low confidence when no context
		AssistantID: assistant.ID,
	}, nil
}

// IndexResponse generates embeddings for an assistant response and stores them for search.
// Called when a response is created or its question/answer is updated.
func (s *RAGService) IndexResponse(ctx context.Context, responseID uint) error {
	resp, err := s.responseRepo.GetByID(ctx, responseID)
	if err != nil {
		return fmt.Errorf("response not found: %w", err)
	}

	// Generate embedding from question + answer for better semantic matching
	inputText := fmt.Sprintf("Q: %s\nA: %s", resp.Question, resp.Answer)
	embedResp, err := s.llmProvider.CreateEmbedding(ctx, llm.EmbeddingRequest{
		Model: "text-embedding-3-small",
		Input: []string{inputText},
	})
	if err != nil {
		applogger.L().Errorf("RAG response embedding: %v", err)
		return fmt.Errorf("response embedding failed: %w", err)
	}

	if len(embedResp.Data) > 0 {
		float32Emb := make([]float32, len(embedResp.Data[0].Embedding))
		for i, v := range embedResp.Data[0].Embedding {
			float32Emb[i] = float32(v)
		}
		resp.Embedding = pgvector.NewVector(float32Emb)
		if err := s.responseRepo.Update(ctx, resp); err != nil {
			applogger.L().Errorf("RAG update response embedding: %v", err)
			return fmt.Errorf("store embedding failed: %w", err)
		}
	}

	applogger.L().Infof("RAG indexed response %d for assistant %d", resp.ID, resp.AssistantID)
	return nil
}