specstream/pkg/extract/markdown.go

// Package extract_md implements the Markdown content extractor using goldmark.
//
// Security model: same as HTMLExtractor — no raw text passthrough.
// Parses markdown AST into typed Section, CodeBlock, List, Table structures.
// Never returns the raw markdown string.

package extract

import (
	"bytes"
	"crypto/sha256"
	"fmt"
	"io"
	"strings"
	"time"

	"github.com/sscoble/specstream/pkg/types"
	"github.com/yuin/goldmark"
	"github.com/yuin/goldmark/ast"
	"github.com/yuin/goldmark/extension"
	east "github.com/yuin/goldmark/extension/ast"
	"github.com/yuin/goldmark/text"
)

type MarkdownExtractor struct{}

func NewMarkdownExtractor() *MarkdownExtractor {
	return &MarkdownExtractor{}
}

func (e *MarkdownExtractor) Detect(r io.ReaderAt, size int64) (bool, error) {
	if size > 1024 {
		size = 1024
	}
	buf := make([]byte, size)
	n, err := r.ReadAt(buf, 0)
	if err != nil && err != io.EOF {
		return false, err
	}
	sample := string(buf[:n])
	hasHeading := len(sample) > 0 && (sample[0] == '#' || strings.Contains(sample, "\n#"))
	hasList := strings.Contains(sample, "- ") || strings.Contains(sample, "* ")
	hasLink := strings.Contains(sample, "](")
	return (hasHeading || hasList) && hasLink, nil
}

func (e *MarkdownExtractor) MIMETypes() []string {
	return []string{"text/markdown", "text/x-markdown"}
}

func (e *MarkdownExtractor) Extract(r io.Reader, req types.ExtractRequest) (*types.ExtractedContent, error) {
	start := time.Now()

	raw, err := io.ReadAll(r)
	if err != nil {
		return nil, fmt.Errorf("read markdown: %w", err)
	}

	hash := sha256.Sum256(raw)
	fingerprint := fmt.Sprintf("%x", hash)

	content := &types.ExtractedContent{
		Source: types.Source{
			MIMEType:  "text/markdown",
			ByteSize:  int64(len(raw)),
			IsTrusted: false,
		},
		ContentType: "markdown",
		ExtractedAt: time.Now(),
		Metadata: types.Metadata{
			ExtractorVersion: "specstream/0.1.0",
			SHA256:           fingerprint,
			StrippedPatterns: make(map[string]int),
			TagCounts:        make(map[string]int),
		},
	}

	// Extract title from first h1
	lines := strings.Split(string(raw), "\n")
	for _, line := range lines {
		trimmed := strings.TrimSpace(line)
		if strings.HasPrefix(trimmed, "# ") {
			content.Title = strings.TrimPrefix(trimmed, "# ")
			break
		}
		if strings.HasPrefix(trimmed, "#") {
			continue
		}
		if trimmed != "" && !strings.HasPrefix(trimmed, "---") {
			if content.Description == "" && len(trimmed) > 10 {
				content.Description = trimmed
				if len(content.Description) > 200 {
					content.Description = content.Description[:200]
				}
			}
			break
		}
	}

	// Parse with goldmark using GFM extensions (tables, etc.)
	md := goldmark.New(
		goldmark.WithExtensions(
			extension.Table,
			extension.Strikethrough,
			extension.TaskList,
		),
	)
	doc := md.Parser().Parse(text.NewReader(raw))

	var currentSection *types.Section

	flushSection := func() {
		if currentSection != nil && len(currentSection.Paragraphs) > 0 {
			content.Sections = append(content.Sections, *currentSection)
		}
	}

	ast.Walk(doc, func(n ast.Node, entering bool) (ast.WalkStatus, error) {
		if !entering {
			return ast.WalkContinue, nil
		}

		switch n.Kind() {
		case ast.KindHeading:
			flushSection()
			h := n.(*ast.Heading)
			text := mdInlineText(h, raw)
			currentSection = &types.Section{
				Heading:    text,
				Level:      h.Level,
				Paragraphs: []string{},
			}

		case ast.KindParagraph:
			// Extract links from paragraphs
			if req.ExtractLinks {
				extractLinksFromNode(n, raw, content)
			}
			text := mdInlineText(n, raw)
			text = sanitizeParagraph(text)
			if text == "" {
				return ast.WalkSkipChildren, nil
			}
			if currentSection == nil {
				currentSection = &types.Section{
					Paragraphs: []string{},
				}
			}
			currentSection.Paragraphs = append(currentSection.Paragraphs, text)

		case ast.KindFencedCodeBlock:
			cb := n.(*ast.FencedCodeBlock)
			lang := string(cb.Language(raw))
			var codeBuf bytes.Buffer
			for child := cb.FirstChild(); child != nil; child = child.NextSibling() {
				if t, ok := child.(*ast.Text); ok {
					codeBuf.Write(t.Segment.Value(raw))
				}
			}
			content.CodeBlocks = append(content.CodeBlocks, types.CodeBlock{
				Language: lang,
				Code:     strings.TrimSpace(codeBuf.String()),
			})

		case ast.KindList:
			l := n.(*ast.List)
			ordered := l.IsOrdered()
			items := []string{}
			for child := l.FirstChild(); child != nil; child = child.NextSibling() {
				if li, ok := child.(*ast.ListItem); ok {
					items = append(items, mdListItemText(li, raw))
					if req.ExtractLinks {
						extractLinksFromNode(li, raw, content)
					}
				}
			}
			if len(items) > 0 {
				content.Lists = append(content.Lists, types.List{
					Ordered: ordered,
					Items:   items,
				})
			}

		case east.KindTable:
			if !req.ExtractTables {
				return ast.WalkSkipChildren, nil
			}
			t := n.(*east.Table)
			table := types.Table{}
			for row := t.FirstChild(); row != nil; row = row.NextSibling() {
				switch row.Kind() {
				case east.KindTableHeader:
					hdr := row.(*east.TableHeader)
					for cell := hdr.FirstChild(); cell != nil; cell = cell.NextSibling() {
						if tc, ok := cell.(*east.TableCell); ok {
							table.Headers = append(table.Headers, mdInlineText(tc, raw))
							if req.ExtractLinks {
								extractLinksFromNode(tc, raw, content)
							}
						}
					}
				case east.KindTableRow:
					tr := row.(*east.TableRow)
					rowData := []string{}
					for cell := tr.FirstChild(); cell != nil; cell = cell.NextSibling() {
						if tc, ok := cell.(*east.TableCell); ok {
							rowData = append(rowData, mdInlineText(tc, raw))
							if req.ExtractLinks {
								extractLinksFromNode(tc, raw, content)
							}
						}
					}
					if len(rowData) > 0 {
						table.Rows = append(table.Rows, rowData)
					}
				}
			}
			if len(table.Rows) > 0 || len(table.Headers) > 0 {
				content.Tables = append(content.Tables, table)
			}
		}

		return ast.WalkContinue, nil
	})
	flushSection()

	content.SectionCnt = len(content.Sections)
	content.CharCount = countChars(content)
	content.Metadata.ExtractionTimeMs = time.Since(start).Milliseconds()

	return content, nil
}

// mdInlineText extracts plain text from inline content nodes.
func mdInlineText(n ast.Node, src []byte) string {
	var buf bytes.Buffer
	for child := n.FirstChild(); child != nil; child = child.NextSibling() {
		switch c := child.(type) {
		case *ast.Text:
			buf.Write(c.Segment.Value(src))
		case *ast.CodeSpan:
			buf.WriteByte('`')
			for cc := c.FirstChild(); cc != nil; cc = cc.NextSibling() {
				if t, ok := cc.(*ast.Text); ok {
					buf.Write(t.Segment.Value(src))
				}
			}
			buf.WriteByte('`')
		case *ast.Link:
			dest := string(c.Destination)
			text := mdInlineText(c, src)
			if text != "" {
				buf.WriteString(text)
				buf.WriteString(" (")
				buf.WriteString(dest)
				buf.WriteString(")")
			} else {
				buf.WriteString(dest)
			}
		case *ast.Image:
			dest := string(c.Destination)
			alt := string(c.Title)
			if alt == "" {
				alt = dest
			}
			buf.WriteString("[image: ")
			buf.WriteString(alt)
			buf.WriteString("]")
		default:
			buf.WriteString(mdInlineText(c, src))
		}
	}
	return strings.TrimSpace(buf.String())
}

// extractLinksFromNode walks a node tree and extracts all Link destinations.
func extractLinksFromNode(n ast.Node, src []byte, content *types.ExtractedContent) {
	for child := n.FirstChild(); child != nil; child = child.NextSibling() {
		if link, ok := child.(*ast.Link); ok {
			dest := string(link.Destination)
			if dest == "" || strings.HasPrefix(dest, "#") {
				continue
			}
			text := mdInlineText(link, src)
			// Deduplicate by URL
			dup := false
			for _, existing := range content.Links {
				if existing.Href == dest {
					dup = true
					break
				}
			}
			if !dup {
				content.Links = append(content.Links, types.Link{
					Text: text,
					Href: dest,
				})
			}
		}
		extractLinksFromNode(child, src, content)
	}
}

// mdListItemText extracts plain text from a ListItem.
func mdListItemText(n ast.Node, src []byte) string {
	var parts []string
	for child := n.FirstChild(); child != nil; child = child.NextSibling() {
		parts = append(parts, mdInlineText(child, src))
	}
	return strings.TrimSpace(strings.Join(parts, " "))
}