package docx

import (
	"archive/zip"
	"encoding/xml"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
)

const (
	MaxDocxSize = 50 * 1024 * 1024 // 50MB
)

// DocumentParser handles DOCX file parsing
type DocumentParser struct {
	filePath string
}

// NewDocumentParser creates a new DOCX parser
func NewDocumentParser(filePath string) *DocumentParser {
	return &DocumentParser{filePath: filePath}
}

// ParseDOCX extracts text from a DOCX file
// Returns (markdown text, extracted images as bytes, error)
func (p *DocumentParser) ParseDOCX() (string, map[string][]byte, error) {
	// Validate file size
	fileInfo, err := os.Stat(p.filePath)
	if err != nil {
		return "", nil, fmt.Errorf("failed to stat file: %w", err)
	}

	if fileInfo.Size() > MaxDocxSize {
		return "", nil, fmt.Errorf("file exceeds 50MB limit: %d bytes", fileInfo.Size())
	}

	// Validate ZIP magic number
	f, err := os.Open(p.filePath)
	if err != nil {
		return "", nil, fmt.Errorf("failed to open file: %w", err)
	}
	defer f.Close()

	magicBytes := make([]byte, 4)
	_, err = f.Read(magicBytes)
	if err != nil {
		return "", nil, fmt.Errorf("failed to read magic bytes: %w", err)
	}

	if string(magicBytes) != "PK\x03\x04" {
		return "", nil, fmt.Errorf("not a valid DOCX file (invalid ZIP header)")
	}

	// Open as ZIP
	f.Seek(0, 0)
	reader, err := zip.NewReader(f, fileInfo.Size())
	if err != nil {
		return "", nil, fmt.Errorf("failed to open ZIP: %w", err)
	}

	// Extract document.xml
	var docText string
	images := make(map[string][]byte)

	for _, file := range reader.File {
		if file.Name == "word/document.xml" {
			rc, err := file.Open()
			if err != nil {
				return "", nil, fmt.Errorf("failed to open document.xml: %w", err)
			}
			defer rc.Close()

			bytes, err := io.ReadAll(rc)
			if err != nil {
				return "", nil, fmt.Errorf("failed to read document.xml: %w", err)
			}

			docText, err = p.extractTextFromXML(bytes)
			if err != nil {
				return "", nil, fmt.Errorf("failed to extract text: %w", err)
			}
		}

		// Extract images from media/
		if strings.HasPrefix(file.Name, "word/media/") {
			rc, err := file.Open()
			if err != nil {
				continue
			}
			defer rc.Close()

			imageBytes, err := io.ReadAll(rc)
			if err != nil {
				continue
			}

			fileName := filepath.Base(file.Name)
			images[fileName] = imageBytes
		}
	}

	return docText, images, nil
}

// extractTextFromXML parses word/document.xml and extracts plain text
// Unsafe XML entity resolution is disabled by default in Go's xml.Decoder
func (p *DocumentParser) extractTextFromXML(data []byte) (string, error) {
	var doc struct {
		XMLName xml.Name
		Body    struct {
			Paragraphs []struct {
				Text     string `xml:"w|t"`
				Runs     []struct {
					Text string `xml:"w|t"`
				} `xml:"w|r"`
			} `xml:"w|p"`
		} `xml:"w|body"`
	}

	decoder := xml.NewDecoder(strings.NewReader(string(data)))
	// XXE protection: Go's xml.Decoder is safe by default (no entity resolution)
	decoder.Strict = false

	err := decoder.Decode(&doc)
	if err != nil && err != io.EOF {
		return "", fmt.Errorf("failed to decode XML: %w", err)
	}

	// Extract text from paragraphs
	var result []string
	for _, p := range doc.Body.Paragraphs {
		var text []string
		for _, run := range p.Runs {
			if run.Text != "" {
				text = append(text, run.Text)
			}
		}
		if len(text) > 0 {
			result = append(result, strings.Join(text, ""))
		}
	}

	return strings.Join(result, "\n"), nil
}