package sitemapgen

import (
	"context"
	"fmt"
	"os"
	"path/filepath"
	"time"
)

// Generator orquestra leitura, escrita XML e checkpoint.
type Generator struct {
	cfg Config
}

func NewGenerator(cfg Config) *Generator {
	return &Generator{cfg: cfg}
}

func (g *Generator) Run(ctx context.Context) error {
	if g.cfg.Reset {
		_ = os.Remove(g.cfg.StatePath)
	}
	var state *State
	if !g.cfg.Reset {
		var err error
		state, err = LoadState(g.cfg.StatePath)
		if err != nil {
			return err
		}
	}
	if state == nil {
		state = &State{
			Version:       1,
			SiteURL:       g.cfg.SiteURL,
			Source:        g.cfg.Source,
			StartedAt:     time.Now().UTC(),
			APIUFProgress: map[string]string{},
		}
	}
	if state.APIUFProgress == nil {
		state.APIUFProgress = map[string]string{}
	}
	if state.Finished && g.cfg.Source == "api" && len(state.APICompletedUFs) >= len(brUFs) {
		fmt.Println("Geração já concluída (todas as UFs processadas). Use -reset para recomeçar.")
		return nil
	}
	if state.Finished && g.cfg.Source == "api" {
		state.Finished = false
	}

	writer, err := NewXMLWriter(g.cfg.SiteURL, g.cfg.OutputDir, g.cfg.URLsPerFile, state)
	if err != nil {
		return err
	}

	store := NewStateStore(g.cfg.StatePath, state)
	var total int64 = state.TotalURLs
	lastCNPJ := state.LastCNPJ

	emit := func(rec Record) error {
		if rec.CNPJ == "" {
			return nil
		}
		if err := writer.Add(rec); err != nil {
			return err
		}
		total++
		lastCNPJ = rec.CNPJ
		if total%int64(g.cfg.BatchSize) == 0 {
			if err := store.Update(func(s *State) {
				s.LastCNPJ = lastCNPJ
				s.TotalURLs = total
				s.CompletedFiles = writer.CompletedFiles()
				s.CurrentFile = writer.CurrentFileName()
				s.URLsInCurrentFile = writer.URLsInCurrentFile()
			}); err != nil {
				return err
			}
			fmt.Printf("checkpoint: %d URLs, último CNPJ %s, arquivo %s (%d)\n",
				total, lastCNPJ, writer.CurrentFileName(), writer.URLsInCurrentFile())
		}
		return nil
	}

	switch g.cfg.Source {
	case "mysql":
		src, err := NewMySQLSource(g.cfg.DBDSN, g.cfg.DBQuery)
		if err != nil {
			return err
		}
		defer src.Close()
		if err := src.Stream(ctx, lastCNPJ, g.cfg.BatchSize, g.cfg.Limit, emit); err != nil {
			_ = g.saveCheckpoint(store, writer, &lastCNPJ, &total)
			return fmt.Errorf("%w — rode de novo sem -reset para retomar", err)
		}
	case "api":
		src := NewAPISource(g.cfg.APIBase, g.cfg.APIKey, g.cfg.Workers)
		onProgress := func(p APIProgress) {
			_ = store.Update(func(s *State) {
				if s.APIUFProgress == nil {
					s.APIUFProgress = map[string]string{}
				}
				if p.Done {
					delete(s.APIUFProgress, p.UF)
					for _, done := range s.APICompletedUFs {
						if done == p.UF {
							return
						}
					}
					s.APICompletedUFs = append(s.APICompletedUFs, p.UF)
					return
				}
				s.APIUFProgress[p.UF] = p.Cursor
			})
		}
		if err := src.Stream(ctx, state.APIUFProgress, state.CompletedUFSet(), g.cfg.Limit, emit, onProgress); err != nil {
			_ = g.saveCheckpoint(store, writer, &lastCNPJ, &total)
			return fmt.Errorf("%w — rode de novo sem -reset para retomar", err)
		}
	default:
		return fmt.Errorf("fonte desconhecida: %s", g.cfg.Source)
	}

	if err := writer.Finish(); err != nil {
		return err
	}
	allUFsDone := g.cfg.Source != "api" || len(state.APICompletedUFs) >= len(brUFs)
	return store.Update(func(s *State) {
		s.LastCNPJ = lastCNPJ
		s.TotalURLs = total
		s.CompletedFiles = writer.CompletedFiles()
		s.CurrentFile = ""
		s.URLsInCurrentFile = 0
		s.Finished = allUFsDone && g.cfg.Limit == 0
	})
}

// OutputIndexPath caminho do índice principal.
func (g *Generator) OutputIndexPath() string {
	return filepath.Join(g.cfg.OutputDir, "sitemap-index.xml")
}

func (g *Generator) saveCheckpoint(store *StateStore, writer *XMLWriter, lastCNPJ *string, total *int64) error {
	return store.Update(func(s *State) {
		s.LastCNPJ = *lastCNPJ
		s.TotalURLs = *total
		s.CompletedFiles = writer.CompletedFiles()
		s.CurrentFile = writer.CurrentFileName()
		s.URLsInCurrentFile = writer.URLsInCurrentFile()
		s.Finished = false
	})
}