package exporter

import (
    "encoding/csv"
    "os"
    "path/filepath"
    "time"
    "github.com/xuri/excelize/v2"
)

type RowWriter interface {
    WriteHeader(cols []string) error
    WriteRow(vals []string) error
    Close() (string, int64, error)
}

type CSVWriter struct {
    f *os.File
    w *csv.Writer
    count int64
}

func NewCSVWriter(dir, name string) (*CSVWriter, error) {
    os.MkdirAll(dir, 0755)
    p := filepath.Join(dir, name+"_"+time.Now().Format("20060102150405")+".csv")
    f, err := os.Create(p)
    if err != nil {
        return nil, err
    }
    return &CSVWriter{f: f, w: csv.NewWriter(f)}, nil
}

func (c *CSVWriter) WriteHeader(cols []string) error {
    if err := c.w.Write(cols); err != nil {
        return err
    }
    c.count++
    return nil
}

func (c *CSVWriter) WriteRow(vals []string) error {
    if err := c.w.Write(vals); err != nil {
        return err
    }
    c.count++
    return nil
}

func (c *CSVWriter) Close() (string, int64, error) {
    c.w.Flush()
    p := c.f.Name()
    info, _ := c.f.Stat()
    c.f.Close()
    return p, info.Size(), nil
}

type XLSXWriter struct {
    f *excelize.File
    sheet string
    row int
}

func NewXLSXWriter(dir, name, sheet string) (*XLSXWriter, string, error) {
    os.MkdirAll(dir, 0755)
    p := filepath.Join(dir, name+"_"+time.Now().Format("20060102150405")+".xlsx")
    f := excelize.NewFile()
    idx, err := f.GetSheetIndex(sheet)
    if err != nil || idx < 0 {
        idx, _ = f.NewSheet(sheet)
        f.SetActiveSheet(idx)
        if sheet != "Sheet1" {
            _ = f.DeleteSheet("Sheet1")
        }
    } else {
        f.SetActiveSheet(idx)
    }
    return &XLSXWriter{f: f, sheet: sheet, row: 1}, p, nil
}

func (x *XLSXWriter) WriteHeader(cols []string) error {
    for i, c := range cols {
        cell := col(i+1) + "1"
        if err := x.f.SetCellValue(x.sheet, cell, c); err != nil {
            return err
        }
    }
    x.row = 2
    return nil
}

func (x *XLSXWriter) WriteRow(vals []string) error {
    r := x.row
    for i, v := range vals {
        cell := col(i+1) + itoa(r)
        if err := x.f.SetCellValue(x.sheet, cell, v); err != nil {
            return err
        }
    }
    x.row++
    return nil
}

func (x *XLSXWriter) Close(path string) (string, int64, error) {
    if err := x.f.SaveAs(path); err != nil {
        return "", 0, err
    }
    info, err := os.Stat(path)
    if err != nil {
        return path, 0, nil
    }
    return path, info.Size(), nil
}

func col(n int) string {
    s := ""
    for n > 0 {
        n--
        s = string(rune('A'+(n%26))) + s
        n /= 26
    }
    return s
}

func itoa(n int) string {
    if n == 0 {
        return "0"
    }
    b := make([]byte, 0, 10)
    m := n
    for m > 0 {
        b = append(b, byte('0'+(m%10)))
        m /= 10
    }
    for i, j := 0, len(b)-1; i < j; i, j = i+1, j-1 {
        b[i], b[j] = b[j], b[i]
    }
    return string(b)
}