package rag

import (
	"context"
	"crypto/sha256"
	"fmt"
	"math"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"time"
)

// Chunk represents a document chunk with its embedding.
type Chunk struct {
	ID        string    `json:"id"`
	DocID     string    `json:"doc_id"`
	DocTitle  string    `json:"doc_title"`
	Content   string    `json:"content"`
	Index     int       `json:"index"`
	Embedding []float64 `json:"-"`
	CreatedAt time.Time `json:"created_at"`
}

// SearchResult represents a retrieved knowledge chunk.
type SearchResult struct {
	Chunk    Chunk   `json:"chunk"`
	Score    float64 `json:"score"`
}

// KnowledgeStore manages document chunks and provides semantic search.
type KnowledgeStore struct {
	mu         sync.RWMutex
	chunks     []Chunk
	embedder   Embedder
	knowledgeDir string
}

// NewKnowledgeStore creates a new knowledge store.
func NewKnowledgeStore(embedder Embedder, knowledgeDir string) *KnowledgeStore {
	if knowledgeDir == "" {
		knowledgeDir = "./data/knowledge"
	}
	return &KnowledgeStore{
		embedder:     embedder,
		knowledgeDir: knowledgeDir,
	}
}

// IngestDirectory scans a directory and indexes all supported files.
func (ks *KnowledgeStore) IngestDirectory(ctx context.Context) (int, error) {
	if _, err := os.Stat(ks.knowledgeDir); os.IsNotExist(err) {
		return 0, nil
	}

	var count int
	err := filepath.Walk(ks.knowledgeDir, func(path string, info os.FileInfo, err error) error {
		if err != nil || info.IsDir() {
			return err
		}
		ext := strings.ToLower(filepath.Ext(path))
		if !isSupportedFile(ext) {
			return nil
		}
		// Skip files that are injected elsewhere (persona.md → system prompt,
		// _index.md → metadata only, not story content).
		base := filepath.Base(path)
		if base == "persona.md" || base == "_index.md" {
			return nil
		}
		n, err := ks.IngestFile(ctx, path)
		if err != nil {
			return fmt.Errorf("ingest %s: %w", path, err)
		}
		count += n
		return nil
	})
	return count, err
}

// IngestFile reads and indexes a single file.
func (ks *KnowledgeStore) IngestFile(ctx context.Context, path string) (int, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return 0, err
	}

	docID := hashString(path)
	title := filepath.Base(path)
	ext := strings.ToLower(filepath.Ext(path))

	var text string
	switch ext {
	case ".md", ".txt", ".go", ".py", ".js", ".ts", ".tsx", ".jsx",
		".json", ".yaml", ".yml", ".toml", ".xml", ".html", ".css",
		".sh", ".bat", ".ps1", ".java", ".rs", ".c", ".cpp", ".h":
		text = string(data)
	default:
		return 0, nil
	}

	chunks := chunkText(text, 1024, 256)
	if len(chunks) == 0 {
		return 0, nil
	}

	texts := make([]string, len(chunks))
	for i, c := range chunks {
		texts[i] = c
	}

	embedding, err := ks.embedder.EmbedBatch(ctx, texts)
	_ = embedding // single embedding for batch — use per-chunk embeddings for accuracy

	var indexed int
	for i, chunk := range chunks {
		chunkID := fmt.Sprintf("%s:%d", docID, i)
		chunkEmbedding, _ := ks.embedder.Embed(ctx, chunk)
		c := Chunk{
			ID:        chunkID,
			DocID:     docID,
			DocTitle:  title,
			Content:   chunk,
			Index:     i,
			Embedding: chunkEmbedding,
			CreatedAt: time.Now(),
		}
		ks.mu.Lock()
		// Replace existing chunks for this doc
		ks.removeDoc(docID)
		ks.chunks = append(ks.chunks, c)
		ks.mu.Unlock()
		indexed++
	}

	return indexed, nil
}

// Search performs semantic search over the knowledge base.
func (ks *KnowledgeStore) Search(ctx context.Context, query string, topK int) ([]SearchResult, error) {
	if topK <= 0 {
		topK = 5
	}

	queryEmbedding, err := ks.embedder.Embed(ctx, query)
	if err != nil {
		return nil, fmt.Errorf("embed query: %w", err)
	}

	ks.mu.RLock()
	defer ks.mu.RUnlock()

	if len(ks.chunks) == 0 {
		return nil, nil
	}

	var results []SearchResult
	for _, chunk := range ks.chunks {
		score := cosineSimilarity(queryEmbedding, chunk.Embedding)
		// Also boost by keyword match
		keywordScore := keywordMatchScore(query, chunk.Content)
		combinedScore := score*0.7 + keywordScore*0.3

		results = append(results, SearchResult{
			Chunk: chunk,
			Score: combinedScore,
		})
	}

	sort.Slice(results, func(i, j int) bool {
		return results[i].Score > results[j].Score
	})

	if len(results) > topK {
		results = results[:topK]
	}

	// Filter out very low relevance
	var filtered []SearchResult
	for _, r := range results {
		if r.Score > 0.01 {
			filtered = append(filtered, r)
		}
	}

	return filtered, nil
}

// Stats returns knowledge base statistics.
func (ks *KnowledgeStore) Stats() map[string]interface{} {
	ks.mu.RLock()
	defer ks.mu.RUnlock()

	docs := make(map[string]int)
	for _, c := range ks.chunks {
		docs[c.DocTitle]++
	}

	return map[string]interface{}{
		"total_chunks":  len(ks.chunks),
		"total_docs":    len(docs),
		"documents":     docs,
		"knowledge_dir": ks.knowledgeDir,
	}
}

func (ks *KnowledgeStore) removeDoc(docID string) {
	filtered := ks.chunks[:0]
	for _, c := range ks.chunks {
		if c.DocID != docID {
			filtered = append(filtered, c)
		}
	}
	ks.chunks = filtered
}

// chunkText splits text into overlapping chunks.
func chunkText(text string, chunkSize, overlap int) []string {
	if len(text) <= chunkSize {
		return []string{text}
	}

	var chunks []string
	runes := []rune(text)
	step := chunkSize - overlap
	if step <= 0 {
		step = chunkSize
	}

	for i := 0; i < len(runes); i += step {
		end := i + chunkSize
		if end > len(runes) {
			end = len(runes)
		}
		chunks = append(chunks, string(runes[i:end]))
		if end == len(runes) {
			break
		}
	}
	return chunks
}

// cosineSimilarity computes cosine similarity between two vectors.
func cosineSimilarity(a, b []float64) float64 {
	if len(a) != len(b) || len(a) == 0 {
		return 0
	}
	var dot, normA, normB float64
	for i := range a {
		dot += a[i] * b[i]
		normA += a[i] * a[i]
		normB += b[i] * b[i]
	}
	if normA == 0 || normB == 0 {
		return 0
	}
	return dot / (math.Sqrt(normA) * math.Sqrt(normB))
}

// keywordMatchScore computes a simple keyword overlap score.
func keywordMatchScore(query, text string) float64 {
	queryLower := strings.ToLower(query)
	textLower := strings.ToLower(text)
	queryWords := strings.Fields(queryLower)
	if len(queryWords) == 0 {
		return 0
	}
	matchCount := 0
	for _, w := range queryWords {
		if len(w) >= 2 && strings.Contains(textLower, w) {
			matchCount++
		}
	}
	return float64(matchCount) / float64(len(queryWords))
}

func hashString(s string) string {
	h := sha256.Sum256([]byte(s))
	return fmt.Sprintf("%x", h[:8])
}

func isSupportedFile(ext string) bool {
	switch ext {
	case ".md", ".txt", ".go", ".py", ".js", ".ts", ".tsx", ".jsx",
		".json", ".yaml", ".yml", ".toml", ".xml", ".html", ".css",
		".sh", ".bat", ".ps1", ".java", ".rs", ".c", ".cpp", ".h":
		return true
	default:
		return false
	}
}
