package sitemapgen

import (
	"context"
	"database/sql"
	"fmt"
	"strings"

	_ "github.com/go-sql-driver/mysql"
)

const defaultMySQLQuery = `
SELECT
  LPAD(CAST(e.cnpj AS CHAR), 14, '0') AS cnpj,
  COALESCE(emp.razao_social, '') AS razao_social,
  NULLIF(TRIM(e.nome_fantasia), '') AS nome_fantasia,
  COALESCE(m.descricao, '') AS municipio_nome,
  COALESCE(e.uf, '') AS uf
FROM estabelecimentos e
LEFT JOIN empresas emp ON emp.cnpj_basico = SUBSTRING(LPAD(CAST(e.cnpj AS CHAR), 14, '0'), 1, 8)
LEFT JOIN municipios m ON m.codigo = e.municipio
WHERE e.situacao_cadastral = 2
  AND LPAD(CAST(e.cnpj AS CHAR), 14, '0') > ?
ORDER BY cnpj
LIMIT ?
`

type MySQLSource struct {
	db    *sql.DB
	query string
}

func NewMySQLSource(dsn, query string) (*MySQLSource, error) {
	db, err := sql.Open("mysql", dsn)
	if err != nil {
		return nil, err
	}
	db.SetMaxOpenConns(4)
	db.SetMaxIdleConns(2)
	if err := db.Ping(); err != nil {
		_ = db.Close()
		return nil, fmt.Errorf("ping mysql: %w", err)
	}
	q := strings.TrimSpace(query)
	if q == "" {
		q = defaultMySQLQuery
	}
	return &MySQLSource{db: db, query: q}, nil
}

func (s *MySQLSource) Close() error {
	return s.db.Close()
}

func (s *MySQLSource) Stream(ctx context.Context, afterCNPJ string, batchSize int, limit int, emit func(Record) error) error {
	last := normalizeCNPJ(afterCNPJ)
	total := 0
	for {
		if err := ctx.Err(); err != nil {
			return err
		}
		rows, err := s.db.QueryContext(ctx, s.query, last, batchSize)
		if err != nil {
			return fmt.Errorf("query: %w", err)
		}
		n := 0
		for rows.Next() {
			var rec Record
			var fantasia sql.NullString
			if err := rows.Scan(&rec.CNPJ, &rec.RazaoSocial, &fantasia, &rec.MunicipioNome, &rec.UF); err != nil {
				rows.Close()
				return err
			}
			rec.CNPJ = normalizeCNPJ(rec.CNPJ)
			if fantasia.Valid {
				rec.NomeFantasia = strings.TrimSpace(fantasia.String)
			}
			last = rec.CNPJ
			n++
			total++
			if err := emit(rec); err != nil {
				rows.Close()
				return err
			}
			if limit > 0 && total >= limit {
				rows.Close()
				return nil
			}
		}
		rows.Close()
		if err := rows.Err(); err != nil {
			return err
		}
		if n == 0 {
			return nil
		}
	}
}

func normalizeCNPJ(s string) string {
	var b strings.Builder
	for _, r := range s {
		if r >= '0' && r <= '9' {
			b.WriteRune(r)
		}
	}
	d := b.String()
	if len(d) > 14 {
		d = d[len(d)-14:]
	}
	for len(d) < 14 {
		d = "0" + d
	}
	return d
}