<commit_before>package scrape

import (
	"compress/gzip"
	"context"
	"encoding/json"
	"encoding/xml"
	"fmt"
	"io"
	"os"
	"path"
	"strconv"
	"strings"
	"time"

	"github.com/slotix/dataflowkit/errs"
	"github.com/slotix/dataflowkit/storage"
	"github.com/spf13/viper"
	"github.com/tealeg/xlsx"
)

const (
	COMMENT_INFO  = "Generated by Dataflow Kit. https://dataflowkit.com"
	GZIP_COMPRESS = "gz"
)

// EncodeToFile save parsed data to specified file.
func EncodeToFile(ctx context.Context, e *encoder, info encodeInfo) ([]byte, error) {
	var w io.Writer
	resultPath := viper.GetString("RESULTS_DIR")
	if _, err := os.Stat(resultPath); os.IsNotExist(err) {
		os.Mkdir(resultPath, 0700)
	}
	timestamp := time.Now().Format("2006-01-02_15:04")
	sFileName := getFilename(resultPath, timestamp, info)
	fo, err := os.OpenFile(sFileName, os.O_CREATE|os.O_WRONLY, 0660)
	if err != nil {
		return nil, err
	}
	defer func() {
		if err := fo.Close(); err != nil {
			panic(err)
		}
	}()

	switch info.compressor {
	case GZIP_COMPRESS:
		gw, _ := gzip.NewWriterLevel(fo, gzip.BestSpeed)
		gw.Name = info.payloadMD5 + "_" + timestamp + "." + info.extension
		gw.Comment = COMMENT_INFO
		w = gw
		defer func() {
			gw.Close()
		}()
	default:
		w = fo
	}

	// close fo on exit and check for its returned error

	var keys *map[int][]int
	if len(info.keys) > 0 {
		keys = info.keys[0]
	}
	err = (*e).encode(ctx, w, info.payloadMD5, keys)
	return []byte(sFileName), err
}

func getFilename(resultPath string, timestamp string, info encodeInfo) (ext string) {
	switch info.compressor {
	case GZIP_COMPRESS:
		ext = GZIP_COMPRESS
	default:
		ext = info.extension
	}
	return path.Join(resultPath, info.payloadMD5+"_"+timestamp+"."+ext)
}

// func EncodeToByteArray(e *encoder, payloadMD5 string, blockMap ...*map[int][]int) ([]byte, error) {
// 	result := ""
// 	buf := bytes.NewBufferString(result)
// 	w := bufio.NewWriter(buf)
// 	var keys *map[int][]int
// 	if len(blockMap) > 0 {
// 		keys = blockMap[0]
// 	}
// 	(*e).encode(w, payloadMD5, keys)
// 	return buf.Bytes(), nil
// }

type encoder interface {
	encode(ctx context.Context, w io.Writer, payloadMD5 string, keys *map[int][]int) error
}

// CSVEncoder transforms parsed data to CSV format.
type CSVEncoder struct {
	partNames []string
	comma     string
}

// JSONEncoder transforms parsed data to JSON format.
type JSONEncoder struct {
	JSONL bool
	//	paginateResults bool
}

// XMLEncoder transforms parsed data to XML format.
type XMLEncoder struct {
}

type XLSXEncoder struct {
	partNames []string
}

func (e JSONEncoder) encode(ctx context.Context, w io.Writer, payloadMD5 string, keys *map[int][]int) error {
	storageType := viper.GetString("STORAGE_TYPE")
	s := storage.NewStore(storageType)
	// make a write buffer
	// if e.paginateResults {
	// 	w.WriteString("[")
	// }
	if !e.JSONL {
		io.WriteString(w, "[")
	}

	reader := newStorageReader(&s, payloadMD5, keys)
	writeComma := false
	for {
		select {
		default:
			block, err := reader.Read()
			if err != nil {
				if err.Error() == errs.EOF {
					if !e.JSONL {
						io.WriteString(w, "]")
					}
					s.Close()
					return nil
				} else if err.Error() == errs.NextPage {
					//next page
					// if e.paginateResults {
					// 	w.WriteString("],[")
					// }
				} else {
					logger.Error(err.Error())
					continue
				}
			}
			if !e.JSONL && writeComma {
				io.WriteString(w, ",")
			}
			blockJSON, err := json.Marshal(block)
			if err != nil {
				logger.Error(err.Error())
			}
			if !writeComma && !e.JSONL {
				writeComma = !writeComma
			}
			w.Write(blockJSON)
			if e.JSONL {
				io.WriteString(w, "\n")
			}
		case <-ctx.Done():
			s.Close()
			return &errs.Cancel{}
		}
	}
}

type storageResultReader struct {
	storage    *storage.Store
	payloadMD5 string
	page       int
	keys       []int
	block      int
	payloadMap map[int][]int
}

func newStorageReader(store *storage.Store, md5Hash string, extractPageKeys *map[int][]int) *storageResultReader {
	reader := &storageResultReader{
		storage:    store,
		payloadMD5: md5Hash,
		payloadMap: make(map[int][]int),
		block:      0,
		page:       0,
	}
	if extractPageKeys != nil {
		reader.payloadMap = *extractPageKeys
	}
	reader.init()
	return reader
}

// have return error
func (r *storageResultReader) init() {
	r.page = 0
	r.block = 0
	if len(r.payloadMap) == 0 {
		keysJSON, err := (*r.storage).Read(storage.Record{
			Type: storage.INTERMEDIATE,
			Key:  r.payloadMD5,
		})
		if err != nil {
			logger.Error(err.Error())
		}
		err = json.Unmarshal(keysJSON, &r.payloadMap)
		if err != nil {
			logger.Error(err.Error())
		}
	}
	for k := range r.payloadMap {
		r.keys = append(r.keys, k)
	}
}

// func (r *storageResultReader) initManualKeys(blocks []int) {
// 	r.keys = blocks
// }

func (r *storageResultReader) Read() (map[string]interface{}, error) {
	//blockMap := make(map[string]interface{})
	var err error
	var nextPage bool
	if r.block >= len(r.payloadMap[r.page]) {
		if r.page+1 < len(r.keys) {
			//achieve next page
			r.page++
			r.block = 0
			nextPage = true
		} else {
			//achieve EOF
			return nil, &errs.ErrStorageResult{Err: errs.EOF}
		}
	}
	blockMap, err := r.getValue()
	if err != nil {
		// have to try get next block value
		r.block++
		return nil, err
	}
	for field, value := range blockMap {
		if strings.Contains(field, "details") {
			details := []map[string]interface{}{}
			detailsReader := newStorageReader(r.storage, value.(string), nil)
			for {
				detailsBlock, detailsErr := detailsReader.Read()
				if detailsErr != nil {
					if detailsErr.Error() == errs.NextPage {

					} else if detailsErr.Error() == errs.EOF {
						break
					} else {
						// in a case of details "no such file or directory" error means, that
						// detail's selector(s) has not be found in a block
						// these can happens when there are few coresponding blocks within a page
						// but only some of them contains wanted selector(s)
						// so just go ahead
						continue
					}
				}
				details = append(details, detailsBlock)
				// we are just breaking here because we got all details recursively
				// if we will continue to read storage, next iteration will returns EOF
				// just to save a time break loop manually
				//break
			}
			if len(details) > 0 {
				if len(details) == 1 {
					blockMap[field] = details[0]
				} else {
					blockMap[field] = details
				}
			}
		}
	}
	r.block++
	if nextPage {
		err = &errs.ErrStorageResult{Err: errs.NextPage}
	}
	return blockMap, err
}

func (r *storageResultReader) getValue() (map[string]interface{}, error) {
	key := fmt.Sprintf("%s-%d-%d", r.payloadMD5, r.page, r.payloadMap[r.page][r.block])
	blockJSON, err := (*r.storage).Read(storage.Record{
		Type: storage.INTERMEDIATE,
		Key:  key,
	})

	if err != nil {
		return nil, err //&errs.ErrStorageResult{Err: fmt.Sprintf(errs.NoKey, key)}
	}
	blockMap := make(map[string]interface{})
	err = json.Unmarshal(blockJSON, &blockMap)
	if err != nil {
		return nil, err
	}
	return blockMap, nil
}

func (e CSVEncoder) encode(ctx context.Context, w io.Writer, payloadMD5 string, keys *map[int][]int) error {
	storageType := viper.GetString("STORAGE_TYPE")
	s := storage.NewStore(storageType)

	//write csv headers
	sString := ""
	for _, headerName := range e.partNames {
		sString += fmt.Sprintf("%s,", headerName)
	}
	sString = strings.TrimSuffix(sString, ",") + "\n"
	_, err := io.WriteString(w, sString)
	if err != nil {
		return err
	}
	//err = w.Flush()
	//if err != nil {
	//	return err
	//}

	reader := newStorageReader(&s, payloadMD5, keys)

	for {
		select {
		default:
			block, err := reader.Read()
			if err != nil {
				if err.Error() == errs.EOF {
					s.Close()
					return nil
				} else if err.Error() == errs.NextPage {
					//next page
				} else {
					logger.Error(err.Error())
					//we have to continue 'cause we still have other records
					continue
				}
			}
			sString = ""
			for _, fieldName := range e.partNames {
				sString += e.formatFieldValue(&block, fieldName)
			}
			sString = strings.TrimSuffix(sString, ",") + "\n"
			_, err = io.WriteString(w, sString)
			if err != nil {
				logger.Error(err.Error())
			}
			//err = w.Flush()
			//if err != nil {
			//	logger.Error(err.Error())
			//}
		case <-ctx.Done():
			s.Close()
			return &errs.Cancel{}
		}
	}

}

func (e CSVEncoder) formatFieldValue(block *map[string]interface{}, fieldName string) string {
	formatedString := ""
	switch v := (*block)[fieldName].(type) {
	case string:
		formatedString = v
		if strings.Contains(formatedString, ",") {
			formatedString = fmt.Sprintf("\"%s\"", formatedString)
		}
	case []string:
		formatedString = strings.Join(v, ";")
	case int:
		formatedString = strconv.FormatInt(int64(v), 10)
	case []int:
		formatedString = intArrayToString(v, ";")
	case []float64:
		formatedString = floatArrayToString(v, ";")
	case float64:
		formatedString = strconv.FormatFloat(v, 'f', -1, 64)
	case nil:
		formatedString = ""
	case []interface{}:
		values := make([]string, len(v))
		for i, value := range v {
			if strings.Contains(value.(string), ",") {
				value = fmt.Sprintf("\"%s\"", value)
			}
			values[i] = fmt.Sprint(value)
		}
		formatedString = strings.Join(values, ";")
	}
	return fmt.Sprintf("%s,", formatedString)
}
<commit_msg>fixed csv format encoder double quotes escaping<commit_after>package scrape

import (
	"compress/gzip"
	"context"
	"encoding/json"
	"encoding/xml"
	"fmt"
	"io"
	"os"
	"path"
	"strconv"
	"strings"
	"time"

	"github.com/slotix/dataflowkit/errs"
	"github.com/slotix/dataflowkit/storage"
	"github.com/spf13/viper"
	"github.com/tealeg/xlsx"
)

const (
	COMMENT_INFO  = "Generated by Dataflow Kit. https://dataflowkit.com"
	GZIP_COMPRESS = "gz"
)

// EncodeToFile save parsed data to specified file.
func EncodeToFile(ctx context.Context, e *encoder, info encodeInfo) ([]byte, error) {
	var w io.Writer
	resultPath := viper.GetString("RESULTS_DIR")
	if _, err := os.Stat(resultPath); os.IsNotExist(err) {
		os.Mkdir(resultPath, 0700)
	}
	timestamp := time.Now().Format("2006-01-02_15:04")
	sFileName := getFilename(resultPath, timestamp, info)
	fo, err := os.OpenFile(sFileName, os.O_CREATE|os.O_WRONLY, 0660)
	if err != nil {
		return nil, err
	}
	defer func() {
		if err := fo.Close(); err != nil {
			panic(err)
		}
	}()

	switch info.compressor {
	case GZIP_COMPRESS:
		gw, _ := gzip.NewWriterLevel(fo, gzip.BestSpeed)
		gw.Name = info.payloadMD5 + "_" + timestamp + "." + info.extension
		gw.Comment = COMMENT_INFO
		w = gw
		defer func() {
			gw.Close()
		}()
	default:
		w = fo
	}

	// close fo on exit and check for its returned error

	var keys *map[int][]int
	if len(info.keys) > 0 {
		keys = info.keys[0]
	}
	err = (*e).encode(ctx, w, info.payloadMD5, keys)
	return []byte(sFileName), err
}

func getFilename(resultPath string, timestamp string, info encodeInfo) (ext string) {
	switch info.compressor {
	case GZIP_COMPRESS:
		ext = GZIP_COMPRESS
	default:
		ext = info.extension
	}
	return path.Join(resultPath, info.payloadMD5+"_"+timestamp+"."+ext)
}

// func EncodeToByteArray(e *encoder, payloadMD5 string, blockMap ...*map[int][]int) ([]byte, error) {
// 	result := ""
// 	buf := bytes.NewBufferString(result)
// 	w := bufio.NewWriter(buf)
// 	var keys *map[int][]int
// 	if len(blockMap) > 0 {
// 		keys = blockMap[0]
// 	}
// 	(*e).encode(w, payloadMD5, keys)
// 	return buf.Bytes(), nil
// }

type encoder interface {
	encode(ctx context.Context, w io.Writer, payloadMD5 string, keys *map[int][]int) error
}

// CSVEncoder transforms parsed data to CSV format.
type CSVEncoder struct {
	partNames []string
	comma     string
}

// JSONEncoder transforms parsed data to JSON format.
type JSONEncoder struct {
	JSONL bool
	//	paginateResults bool
}

// XMLEncoder transforms parsed data to XML format.
type XMLEncoder struct {
}

type XLSXEncoder struct {
	partNames []string
}

func (e JSONEncoder) encode(ctx context.Context, w io.Writer, payloadMD5 string, keys *map[int][]int) error {
	storageType := viper.GetString("STORAGE_TYPE")
	s := storage.NewStore(storageType)
	// make a write buffer
	// if e.paginateResults {
	// 	w.WriteString("[")
	// }
	if !e.JSONL {
		io.WriteString(w, "[")
	}

	reader := newStorageReader(&s, payloadMD5, keys)
	writeComma := false
	for {
		select {
		default:
			block, err := reader.Read()
			if err != nil {
				if err.Error() == errs.EOF {
					if !e.JSONL {
						io.WriteString(w, "]")
					}
					s.Close()
					return nil
				} else if err.Error() == errs.NextPage {
					//next page
					// if e.paginateResults {
					// 	w.WriteString("],[")
					// }
				} else {
					logger.Error(err.Error())
					continue
				}
			}
			if !e.JSONL && writeComma {
				io.WriteString(w, ",")
			}
			blockJSON, err := json.Marshal(block)
			if err != nil {
				logger.Error(err.Error())
			}
			if !writeComma && !e.JSONL {
				writeComma = !writeComma
			}
			w.Write(blockJSON)
			if e.JSONL {
				io.WriteString(w, "\n")
			}
		case <-ctx.Done():
			s.Close()
			return &errs.Cancel{}
		}
	}
}

type storageResultReader struct {
	storage    *storage.Store
	payloadMD5 string
	page       int
	keys       []int
	block      int
	payloadMap map[int][]int
}

func newStorageReader(store *storage.Store, md5Hash string, extractPageKeys *map[int][]int) *storageResultReader {
	reader := &storageResultReader{
		storage:    store,
		payloadMD5: md5Hash,
		payloadMap: make(map[int][]int),
		block:      0,
		page:       0,
	}
	if extractPageKeys != nil {
		reader.payloadMap = *extractPageKeys
	}
	reader.init()
	return reader
}

// have return error
func (r *storageResultReader) init() {
	r.page = 0
	r.block = 0
	if len(r.payloadMap) == 0 {
		keysJSON, err := (*r.storage).Read(storage.Record{
			Type: storage.INTERMEDIATE,
			Key:  r.payloadMD5,
		})
		if err != nil {
			logger.Error(err.Error())
		}
		err = json.Unmarshal(keysJSON, &r.payloadMap)
		if err != nil {
			logger.Error(err.Error())
		}
	}
	for k := range r.payloadMap {
		r.keys = append(r.keys, k)
	}
}

// func (r *storageResultReader) initManualKeys(blocks []int) {
// 	r.keys = blocks
// }

func (r *storageResultReader) Read() (map[string]interface{}, error) {
	//blockMap := make(map[string]interface{})
	var err error
	var nextPage bool
	if r.block >= len(r.payloadMap[r.page]) {
		if r.page+1 < len(r.keys) {
			//achieve next page
			r.page++
			r.block = 0
			nextPage = true
		} else {
			//achieve EOF
			return nil, &errs.ErrStorageResult{Err: errs.EOF}
		}
	}
	blockMap, err := r.getValue()
	if err != nil {
		// have to try get next block value
		r.block++
		return nil, err
	}
	for field, value := range blockMap {
		if strings.Contains(field, "details") {
			details := []map[string]interface{}{}
			detailsReader := newStorageReader(r.storage, value.(string), nil)
			for {
				detailsBlock, detailsErr := detailsReader.Read()
				if detailsErr != nil {
					if detailsErr.Error() == errs.NextPage {

					} else if detailsErr.Error() == errs.EOF {
						break
					} else {
						// in a case of details "no such file or directory" error means, that
						// detail's selector(s) has not be found in a block
						// these can happens when there are few coresponding blocks within a page
						// but only some of them contains wanted selector(s)
						// so just go ahead
						continue
					}
				}
				details = append(details, detailsBlock)
				// we are just breaking here because we got all details recursively
				// if we will continue to read storage, next iteration will returns EOF
				// just to save a time break loop manually
				//break
			}
			if len(details) > 0 {
				if len(details) == 1 {
					blockMap[field] = details[0]
				} else {
					blockMap[field] = details
				}
			}
		}
	}
	r.block++
	if nextPage {
		err = &errs.ErrStorageResult{Err: errs.NextPage}
	}
	return blockMap, err
}

func (r *storageResultReader) getValue() (map[string]interface{}, error) {
	key := fmt.Sprintf("%s-%d-%d", r.payloadMD5, r.page, r.payloadMap[r.page][r.block])
	blockJSON, err := (*r.storage).Read(storage.Record{
		Type: storage.INTERMEDIATE,
		Key:  key,
	})

	if err != nil {
		return nil, err //&errs.ErrStorageResult{Err: fmt.Sprintf(errs.NoKey, key)}
	}
	blockMap := make(map[string]interface{})
	err = json.Unmarshal(blockJSON, &blockMap)
	if err != nil {
		return nil, err
	}
	return blockMap, nil
}

func (e CSVEncoder) encode(ctx context.Context, w io.Writer, payloadMD5 string, keys *map[int][]int) error {
	storageType := viper.GetString("STORAGE_TYPE")
	s := storage.NewStore(storageType)

	//write csv headers
	sString := ""
	for _, headerName := range e.partNames {
		sString += fmt.Sprintf("%s,", headerName)
	}
	sString = strings.TrimSuffix(sString, ",") + "\n"
	_, err := io.WriteString(w, sString)
	if err != nil {
		return err
	}
	//err = w.Flush()
	//if err != nil {
	//	return err
	//}

	reader := newStorageReader(&s, payloadMD5, keys)

	for {
		select {
		default:
			block, err := reader.Read()
			if err != nil {
				if err.Error() == errs.EOF {
					s.Close()
					return nil
				} else if err.Error() == errs.NextPage {
					//next page
				} else {
					logger.Error(err.Error())
					//we have to continue 'cause we still have other records
					continue
				}
			}
			sString = ""
			for _, fieldName := range e.partNames {
				sString += e.formatFieldValue(&block, fieldName)
			}
			sString = strings.TrimSuffix(sString, ",") + "\n"
			_, err = io.WriteString(w, sString)
			if err != nil {
				logger.Error(err.Error())
			}
			//err = w.Flush()
			//if err != nil {
			//	logger.Error(err.Error())
			//}
		case <-ctx.Done():
			s.Close()
			return &errs.Cancel{}
		}
	}

}

func (e CSVEncoder) formatFieldValue(block *map[string]interface{}, fieldName string) string {
	formatedString := ""
	switch v := (*block)[fieldName].(type) {
	case string:
		formatedString = v
		if strings.Contains(formatedString, ",") {
			formatedString = fmt.Sprintf("\"%s\"", formatedString)
		}
	case []string:
		formatedString = strings.Join(v, ";")
	case int:
		formatedString = strconv.FormatInt(int64(v), 10)
	case []int:
		formatedString = intArrayToString(v, ";")
	case []float64:
		formatedString = floatArrayToString(v, ";")
	case float64:
		formatedString = strconv.FormatFloat(v, 'f', -1, 64)
	case nil:
		formatedString = ""
	case []interface{}:
		values := make([]string, len(v))
		for i, value := range v {
			if strings.Contains(value.(string), `"`) {
				value = strings.Replace(value.(string), `"`, `""`, -1)
			}
			values[i] = fmt.Sprint(value)
		}
		formatedString = strings.Join(values, ";")
		if strings.Contains(formatedString, ",") || strings.Contains(formatedString, "\n") {
			formatedString = fmt.Sprintf("\"%s\"", formatedString)
		}
	}
	return fmt.Sprintf("%s,", formatedString)
}

{'fraction_non_alphanumeric': 0.15346825439205425, 'fraction_numerical': 0.006781652159982912, 'mean_word_length': 3.1396993810786915, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 11, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 2}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3555437', 'n_tokens_mistral': 7619, 'n_tokens_neox': 6747, 'n_words': 1641}