package main

import (
	"archive/zip"
	"bytes"
	"context"
	"crypto/sha256"
	"database/sql"
	"encoding/hex"
	"encoding/xml"
	"errors"
	"flag"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	_ "modernc.org/sqlite"
)

type Table struct {
	Index int
	Rows  []Row
}

type Row struct {
	Index int
	Cells []Cell
}

type Cell struct {
	RowIndex  int
	CellIndex int
	ColIndex  int
	ColSpan   int
	VMerge    string
	Text      string
}

type Field struct {
	RecordIndex int
	RowIndex    int
	ColIndex    int
	Key         string
	Value       string
	SourceKind  string
}

func main() {
	if err := run(os.Args[1:]); err != nil {
		fmt.Fprintln(os.Stderr, "error:", err)
		os.Exit(1)
	}
}

func run(args []string) error {
	fs := flag.NewFlagSet("wordtable2sqlite", flag.ContinueOnError)
	dbPath := fs.String("db", "word_tables.sqlite", "SQLite database path")
	reset := fs.Bool("reset", false, "remove the target database before importing")
	quiet := fs.Bool("quiet", false, "suppress per-file summary output")
	if err := fs.Parse(args); err != nil {
		return err
	}
	docxPaths := fs.Args()
	if len(docxPaths) == 0 {
		return errors.New("usage: wordtable2sqlite -db output.sqlite [-reset] file1.docx [file2.docx ...]")
	}

	if *reset {
		if err := os.Remove(*dbPath); err != nil && !errors.Is(err, os.ErrNotExist) {
			return fmt.Errorf("reset database: %w", err)
		}
	}
	if dir := filepath.Dir(*dbPath); dir != "." {
		if err := os.MkdirAll(dir, 0o755); err != nil {
			return fmt.Errorf("create database directory: %w", err)
		}
	}

	db, err := sql.Open("sqlite", *dbPath)
	if err != nil {
		return fmt.Errorf("open sqlite database: %w", err)
	}
	defer db.Close()

	ctx := context.Background()
	for _, docxPath := range docxPaths {
		before, err := dbStats(ctx, db)
		if err != nil {
			return err
		}
		if err := importDocument(ctx, db, docxPath); err != nil {
			return err
		}
		after, err := dbStats(ctx, db)
		if err != nil {
			return err
		}
		if !*quiet {
			fmt.Printf("imported %s: %d tables, %d cells, %d fields\n",
				docxPath,
				after.Tables-before.Tables,
				after.Cells-before.Cells,
				after.Fields-before.Fields,
			)
		}
	}
	if !*quiet {
		fmt.Printf("database: %s\n", *dbPath)
	}
	return nil
}

type stats struct {
	Tables int
	Cells  int
	Fields int
}

func dbStats(ctx context.Context, db *sql.DB) (stats, error) {
	if err := ensureSchema(ctx, db); err != nil {
		return stats{}, err
	}
	var s stats
	for _, item := range []struct {
		query string
		dst   *int
	}{
		{`SELECT COUNT(*) FROM tables`, &s.Tables},
		{`SELECT COUNT(*) FROM cells`, &s.Cells},
		{`SELECT COUNT(*) FROM fields`, &s.Fields},
	} {
		if err := db.QueryRowContext(ctx, item.query).Scan(item.dst); err != nil {
			return stats{}, err
		}
	}
	return s, nil
}

func importDocument(ctx context.Context, db *sql.DB, docxPath string) error {
	if err := ensureSchema(ctx, db); err != nil {
		return err
	}
	tables, err := readDocxTables(docxPath)
	if err != nil {
		return err
	}
	sum, err := fileSHA256(docxPath)
	if err != nil {
		return err
	}

	tx, err := db.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	defer tx.Rollback()

	res, err := tx.ExecContext(ctx,
		`INSERT INTO documents(path, filename, sha256, imported_at) VALUES (?, ?, ?, ?)`,
		docxPath,
		filepath.Base(docxPath),
		sum,
		time.Now().Format(time.RFC3339),
	)
	if err != nil {
		return fmt.Errorf("insert document: %w", err)
	}
	documentID, err := res.LastInsertId()
	if err != nil {
		return err
	}

	for tableIndex, table := range tables {
		cellCount := 0
		for _, row := range table.Rows {
			cellCount += len(row.Cells)
		}
		res, err := tx.ExecContext(ctx,
			`INSERT INTO tables(document_id, table_index, row_count, cell_count) VALUES (?, ?, ?, ?)`,
			documentID,
			tableIndex,
			len(table.Rows),
			cellCount,
		)
		if err != nil {
			return fmt.Errorf("insert table %d: %w", tableIndex, err)
		}
		tableID, err := res.LastInsertId()
		if err != nil {
			return err
		}

		for rowIndex, row := range table.Rows {
			for cellIndex, cell := range row.Cells {
				if _, err := tx.ExecContext(ctx,
					`INSERT INTO cells(document_id, table_id, table_index, row_index, cell_index, col_index, col_span, v_merge, text)
					 VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)`,
					documentID,
					tableID,
					tableIndex,
					rowIndex,
					cellIndex,
					cell.ColIndex,
					cell.ColSpan,
					cell.VMerge,
					cell.Text,
				); err != nil {
					return fmt.Errorf("insert cell table=%d row=%d cell=%d: %w", tableIndex, rowIndex, cellIndex, err)
				}
			}
		}

		for _, field := range inferFields(table) {
			if _, err := tx.ExecContext(ctx,
				`INSERT INTO fields(document_id, table_id, table_index, record_index, row_index, col_index, key, value, source_kind)
				 VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)`,
				documentID,
				tableID,
				tableIndex,
				field.RecordIndex,
				field.RowIndex,
				field.ColIndex,
				field.Key,
				field.Value,
				field.SourceKind,
			); err != nil {
				return fmt.Errorf("insert field %q: %w", field.Key, err)
			}
		}
	}

	return tx.Commit()
}

func ensureSchema(ctx context.Context, db execer) error {
	statements := []string{
		`PRAGMA foreign_keys = ON`,
		`CREATE TABLE IF NOT EXISTS documents (
			id INTEGER PRIMARY KEY AUTOINCREMENT,
			path TEXT NOT NULL,
			filename TEXT NOT NULL,
			sha256 TEXT NOT NULL,
			imported_at TEXT NOT NULL
		)`,
		`CREATE TABLE IF NOT EXISTS tables (
			id INTEGER PRIMARY KEY AUTOINCREMENT,
			document_id INTEGER NOT NULL REFERENCES documents(id) ON DELETE CASCADE,
			table_index INTEGER NOT NULL,
			row_count INTEGER NOT NULL,
			cell_count INTEGER NOT NULL
		)`,
		`CREATE TABLE IF NOT EXISTS cells (
			id INTEGER PRIMARY KEY AUTOINCREMENT,
			document_id INTEGER NOT NULL REFERENCES documents(id) ON DELETE CASCADE,
			table_id INTEGER NOT NULL REFERENCES tables(id) ON DELETE CASCADE,
			table_index INTEGER NOT NULL,
			row_index INTEGER NOT NULL,
			cell_index INTEGER NOT NULL,
			col_index INTEGER NOT NULL,
			col_span INTEGER NOT NULL,
			v_merge TEXT NOT NULL DEFAULT '',
			text TEXT NOT NULL
		)`,
		`CREATE TABLE IF NOT EXISTS fields (
			id INTEGER PRIMARY KEY AUTOINCREMENT,
			document_id INTEGER NOT NULL REFERENCES documents(id) ON DELETE CASCADE,
			table_id INTEGER NOT NULL REFERENCES tables(id) ON DELETE CASCADE,
			table_index INTEGER NOT NULL,
			record_index INTEGER NOT NULL,
			row_index INTEGER NOT NULL,
			col_index INTEGER NOT NULL,
			key TEXT NOT NULL,
			value TEXT NOT NULL,
			source_kind TEXT NOT NULL
		)`,
		`CREATE INDEX IF NOT EXISTS idx_cells_doc_table ON cells(document_id, table_index, row_index, cell_index)`,
		`CREATE INDEX IF NOT EXISTS idx_fields_doc_key ON fields(document_id, key)`,
	}
	for _, statement := range statements {
		if _, err := db.ExecContext(ctx, statement); err != nil {
			return fmt.Errorf("schema statement failed: %w", err)
		}
	}
	return nil
}

type execer interface {
	ExecContext(context.Context, string, ...any) (sql.Result, error)
}

func readDocxTables(path string) ([]Table, error) {
	zr, err := zip.OpenReader(path)
	if err != nil {
		return nil, fmt.Errorf("open docx zip: %w", err)
	}
	defer zr.Close()

	for _, file := range zr.File {
		if file.Name != "word/document.xml" {
			continue
		}
		rc, err := file.Open()
		if err != nil {
			return nil, err
		}
		defer rc.Close()
		data, err := io.ReadAll(rc)
		if err != nil {
			return nil, err
		}
		return parseDocumentXML(data)
	}
	return nil, errors.New("word/document.xml not found in docx")
}

func parseDocumentXML(data []byte) ([]Table, error) {
	decoder := xml.NewDecoder(bytes.NewReader(data))
	var tables []Table
	for {
		token, err := decoder.Token()
		if errors.Is(err, io.EOF) {
			return tables, nil
		}
		if err != nil {
			return nil, err
		}
		start, ok := token.(xml.StartElement)
		if !ok || start.Name.Local != "tbl" {
			continue
		}
		table, err := parseTable(decoder)
		if err != nil {
			return nil, err
		}
		table.Index = len(tables)
		tables = append(tables, table)
	}
}

func parseTable(decoder *xml.Decoder) (Table, error) {
	var table Table
	for {
		token, err := decoder.Token()
		if err != nil {
			return table, err
		}
		switch t := token.(type) {
		case xml.StartElement:
			if t.Name.Local == "tr" {
				row, err := parseRow(decoder, len(table.Rows))
				if err != nil {
					return table, err
				}
				table.Rows = append(table.Rows, row)
			}
		case xml.EndElement:
			if t.Name.Local == "tbl" {
				return table, nil
			}
		}
	}
}

func parseRow(decoder *xml.Decoder, rowIndex int) (Row, error) {
	row := Row{Index: rowIndex}
	logicalCol := 0
	for {
		token, err := decoder.Token()
		if err != nil {
			return row, err
		}
		switch t := token.(type) {
		case xml.StartElement:
			if t.Name.Local == "tc" {
				cell, err := parseCell(decoder)
				if err != nil {
					return row, err
				}
				cell.RowIndex = rowIndex
				cell.CellIndex = len(row.Cells)
				cell.ColIndex = logicalCol
				if cell.ColSpan < 1 {
					cell.ColSpan = 1
				}
				logicalCol += cell.ColSpan
				row.Cells = append(row.Cells, cell)
			}
		case xml.EndElement:
			if t.Name.Local == "tr" {
				return row, nil
			}
		}
	}
}

func parseCell(decoder *xml.Decoder) (Cell, error) {
	cell := Cell{ColSpan: 1}
	var b strings.Builder
	inText := false
	tcDepth := 1
	for {
		token, err := decoder.Token()
		if err != nil {
			return cell, err
		}
		switch t := token.(type) {
		case xml.StartElement:
			switch t.Name.Local {
			case "tc":
				tcDepth++
			case "gridSpan":
				if tcDepth == 1 {
					if span, err := strconv.Atoi(attr(t, "val")); err == nil && span > 0 {
						cell.ColSpan = span
					}
				}
			case "vMerge":
				if tcDepth == 1 {
					cell.VMerge = attr(t, "val")
					if cell.VMerge == "" {
						cell.VMerge = "continue"
					}
				}
			case "t":
				inText = true
			case "tab":
				b.WriteByte('\t')
			case "br", "cr":
				b.WriteByte('\n')
			}
		case xml.CharData:
			if inText {
				b.Write([]byte(t))
			}
		case xml.EndElement:
			switch t.Name.Local {
			case "t":
				inText = false
			case "p":
				if b.Len() > 0 && !strings.HasSuffix(b.String(), "\n") {
					b.WriteByte('\n')
				}
			case "tc":
				tcDepth--
				if tcDepth == 0 {
					cell.Text = normalizeCellText(b.String())
					return cell, nil
				}
			}
		}
	}
}

func attr(start xml.StartElement, local string) string {
	for _, a := range start.Attr {
		if a.Name.Local == local {
			return a.Value
		}
	}
	return ""
}

func normalizeCellText(s string) string {
	s = strings.ReplaceAll(s, "\r\n", "\n")
	s = strings.ReplaceAll(s, "\r", "\n")
	lines := strings.Split(s, "\n")
	out := make([]string, 0, len(lines))
	for _, line := range lines {
		line = strings.TrimSpace(strings.Join(strings.Fields(line), " "))
		if line != "" {
			out = append(out, line)
		}
	}
	return strings.Join(out, "\n")
}

func inferFields(table Table) []Field {
	headerIndex := detectHeaderRow(table)
	if headerIndex >= 0 {
		return inferHeaderFields(table, headerIndex)
	}
	return inferKeyValueFields(table)
}

func detectHeaderRow(table Table) int {
	for i := 0; i < len(table.Rows)-1; i++ {
		row := table.Rows[i]
		if isTitleRow(row) {
			continue
		}
		if nonEmptyCellCount(row) >= 2 && nonEmptyCellCount(table.Rows[i+1]) >= 2 {
			return i
		}
	}
	return -1
}

func isTitleRow(row Row) bool {
	nonEmpty := nonEmptyCellCount(row)
	if nonEmpty != 1 || len(row.Cells) == 0 {
		return false
	}
	return len(row.Cells) == 1 || row.Cells[0].ColSpan > 1
}

func nonEmptyCellCount(row Row) int {
	count := 0
	for _, cell := range row.Cells {
		if strings.TrimSpace(cell.Text) != "" {
			count++
		}
	}
	return count
}

func inferHeaderFields(table Table, headerIndex int) []Field {
	headers := table.Rows[headerIndex].Cells
	var fields []Field
	recordIndex := 0
	for rowIndex := headerIndex + 1; rowIndex < len(table.Rows); rowIndex++ {
		row := table.Rows[rowIndex]
		if nonEmptyCellCount(row) == 0 {
			continue
		}
		for i, header := range headers {
			if i >= len(row.Cells) {
				continue
			}
			key := cleanFieldKey(header.Text)
			if key == "" {
				continue
			}
			value := strings.TrimSpace(row.Cells[i].Text)
			fields = append(fields, Field{
				RecordIndex: recordIndex,
				RowIndex:    rowIndex,
				ColIndex:    row.Cells[i].ColIndex,
				Key:         key,
				Value:       value,
				SourceKind:  "header_row",
			})
		}
		recordIndex++
	}
	return fields
}

func inferKeyValueFields(table Table) []Field {
	var fields []Field
	recordIndex := 0
	for rowIndex, row := range table.Rows {
		cells := nonEmptyCells(row)
		if len(cells) < 2 || len(cells)%2 != 0 {
			continue
		}
		for i := 0; i < len(cells); i += 2 {
			key := cleanFieldKey(cells[i].Text)
			if key == "" {
				continue
			}
			fields = append(fields, Field{
				RecordIndex: recordIndex,
				RowIndex:    rowIndex,
				ColIndex:    cells[i+1].ColIndex,
				Key:         key,
				Value:       cells[i+1].Text,
				SourceKind:  "key_value",
			})
		}
		recordIndex++
	}
	return fields
}

func nonEmptyCells(row Row) []Cell {
	out := make([]Cell, 0, len(row.Cells))
	for _, cell := range row.Cells {
		if strings.TrimSpace(cell.Text) != "" {
			out = append(out, cell)
		}
	}
	return out
}

func cleanFieldKey(s string) string {
	s = strings.TrimSpace(s)
	s = strings.TrimRight(s, "：: \t\r\n")
	return strings.TrimSpace(s)
}

func fileSHA256(path string) (string, error) {
	f, err := os.Open(path)
	if err != nil {
		return "", err
	}
	defer f.Close()
	h := sha256.New()
	if _, err := io.Copy(h, f); err != nil {
		return "", err
	}
	return hex.EncodeToString(h.Sum(nil)), nil
}