package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"strings"
	"time"
)

const maxLogSnippetLen = 256

// StatementRecord is one JSON line from statement source files.
type StatementRecord struct {
	Begin              int64  `json:"begin"`
	End                int64  `json:"end"`
	ServerlessTenantID string `json:"serverless_tenant_id"`
	Raw                json.RawMessage
}

type StatementStreamParser struct {
	reader *bufio.Reader
}

func NewStatementStreamParser(reader io.Reader) *StatementStreamParser {
	return &StatementStreamParser{reader: bufio.NewReader(reader)}
}

func (p *StatementStreamParser) ReadNext() (*StatementRecord, error) {
	for {
		line, err := p.reader.ReadString('\n')
		if err != nil {
			if err == io.EOF {
				line = strings.TrimSpace(line)
				if line == "" {
					return nil, nil
				}
				return parseStatementLine(line)
			}
			return nil, fmt.Errorf("read statement line: %w", err)
		}

		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}
		return parseStatementLine(line)
	}
}

func parseStatementLine(line string) (*StatementRecord, error) {
	var rec StatementRecord
	if err := json.Unmarshal([]byte(line), &rec); err != nil {
		return nil, fmt.Errorf("invalid statement json: %w; line=%s", err, truncateForLog(line, maxLogSnippetLen))
	}
	rec.Raw = json.RawMessage(line)
	return &rec, nil
}

// statementCSVColumnsBase is the stable CSV schema for sqlstatement sync.
// Nested objects (backoff_types, auth_users) are serialized as JSON strings.
var statementCSVColumnsBase = []string{
	"begin", "end", "schema_name", "digest", "plan_digest", "stmt_type",
	"normalized_sql", "table_names", "is_internal", "sample_sql", "charset",
	"collation", "prev_sql", "sample_plan", "sample_binary_plan", "plan_hint",
	"index_names", "exec_count", "sum_errors", "sum_warnings", "sum_latency",
	"max_latency", "min_latency", "sum_parse_latency", "max_parse_latency",
	"sum_compile_latency", "max_compile_latency", "sum_num_cop_tasks",
	"max_cop_process_time", "max_cop_process_address", "max_cop_wait_time",
	"max_cop_wait_address", "sum_process_time", "max_process_time",
	"sum_wait_time", "max_wait_time", "sum_backoff_time", "max_backoff_time",
	"sum_total_keys", "max_total_keys", "sum_processed_keys", "max_processed_keys",
	"sum_rocksdb_delete_skipped_count", "max_rocksdb_delete_skipped_count",
	"sum_rocksdb_key_skipped_count", "max_rocksdb_key_skipped_count",
	"sum_rocksdb_block_cache_hit_count", "max_rocksdb_block_cache_hit_count",
	"sum_rocksdb_block_read_count", "max_rocksdb_block_read_count",
	"sum_rocksdb_block_read_byte", "max_rocksdb_block_read_byte",
	"sum_ia_read_segment_count", "max_ia_read_segment_count",
	"sum_ia_remote_read_segment_size", "max_ia_remote_read_segment_size",
	"sum_ia_remote_read_segment_wait_time", "max_ia_remote_read_segment_wait_time",
	"commit_count", "sum_get_commit_ts_time", "max_get_commit_ts_time",
	"sum_prewrite_time", "max_prewrite_time", "sum_commit_time", "max_commit_time",
	"sum_local_latch_time", "max_local_latch_time", "sum_commit_backoff_time",
	"max_commit_backoff_time", "sum_resolve_lock_time", "max_resolve_lock_time",
	"sum_write_keys", "max_write_keys", "sum_write_size", "max_write_size",
	"sum_prewrite_region_num", "max_prewrite_region_num", "sum_txn_retry",
	"max_txn_retry", "sum_backoff_times", "backoff_types", "auth_users",
	"sum_mem", "max_mem", "sum_disk", "max_disk", "sum_affected_rows",
	"sum_kv_total", "sum_pd_total", "sum_backoff_total", "sum_write_sql_resp_total",
	"sum_tidb_cpu", "sum_tikv_cpu", "sum_result_rows", "max_result_rows",
	"min_result_rows", "prepared", "first_seen", "last_seen", "plan_in_cache",
	"plan_cache_hits", "plan_in_binding", "exec_retry_count", "exec_retry_time",
	"keyspace_name", "keyspace_id", "serverless_tenant_id", "serverless_project_id",
	"serverless_cluster_id", "resource_group_name", "sum_rru", "sum_wru",
	"sum_ru_wait_duration", "max_rru", "max_wru", "max_ru_wait_duration",
	"plan_cache_unqualified_count", "plan_cache_unqualified_last_reason",
}

// statementRuV2CSVColumns are only written when -enable-ru-v2 is true.
var statementRuV2CSVColumns = []string{
	"sum_ruv2",
	"max_ruv2",
	"AVG_REQUEST_UNIT_V2",
	"MAX_REQUEST_UNIT_V2",
}

func statementCSVHeader(includeRuV2 bool) []string {
	if !includeRuV2 {
		out := make([]string, len(statementCSVColumnsBase))
		copy(out, statementCSVColumnsBase)
		return out
	}
	out := make([]string, 0, len(statementCSVColumnsBase)+len(statementRuV2CSVColumns))
	out = append(out, statementCSVColumnsBase...)
	return append(out, statementRuV2CSVColumns...)
}

// statementJSONToCSVRow maps one statement JSON object onto the CSV schema.
func statementJSONToCSVRow(raw json.RawMessage, includeRuV2 bool) ([]string, error) {
	var obj map[string]json.RawMessage
	if err := json.Unmarshal(raw, &obj); err != nil {
		return nil, fmt.Errorf("unmarshal statement for csv: %w", err)
	}
	lookup := make(map[string]json.RawMessage, len(obj))
	for k, v := range obj {
		lookup[strings.ToLower(k)] = v
	}

	cols := statementCSVHeader(includeRuV2)
	row := make([]string, len(cols))
	for i, col := range cols {
		v, ok := lookup[strings.ToLower(col)]
		if !ok {
			row[i] = ""
			continue
		}
		row[i] = jsonValueToCSVCell(v)
	}
	return row, nil
}

func jsonValueToCSVCell(v json.RawMessage) string {
	b := bytes.TrimSpace(v)
	if len(b) == 0 || bytes.Equal(b, []byte("null")) {
		return ""
	}
	switch b[0] {
	case '"':
		var s string
		if err := json.Unmarshal(b, &s); err != nil {
			return string(b)
		}
		return s
	case '{', '[':
		return string(b)
	default:
		return string(b)
	}
}

func (r *StatementRecord) eventTime() time.Time {
	if r.End > 0 {
		return time.Unix(r.End, 0).UTC()
	}
	if r.Begin > 0 {
		return time.Unix(r.Begin, 0).UTC()
	}
	return time.Time{}
}

func truncateForLog(s string, max int) string {
	if max <= 0 || len(s) <= max {
		return s
	}
	return s[:max] + "...(truncated)"
}