package certificatetransparency

import (
	"context"
	"errors"
	"fmt"
	"io"
	"log"
	"net/http"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/google/trillian/client/backoff"

	"github.com/d-Rickyy-b/certstream-server-go/internal/config"
	"github.com/d-Rickyy-b/certstream-server-go/internal/metrics"
	"github.com/d-Rickyy-b/certstream-server-go/internal/models"
	"github.com/d-Rickyy-b/certstream-server-go/internal/web"

	ct "github.com/google/certificate-transparency-go"
	"github.com/google/certificate-transparency-go/client"
	"github.com/google/certificate-transparency-go/jsonclient"
	"github.com/google/certificate-transparency-go/loglist3"
	"github.com/google/certificate-transparency-go/scanner"
)

var (
	errCreatingClient    = errors.New("failed to create JSON client")
	errFetchingSTHFailed = errors.New("failed to fetch STH")
	userAgent            = fmt.Sprintf("Certstream Server v%s (github.com/d-Rickyy-b/certstream-server-go)", config.Version)
)

// Watcher describes a component that watches for new certificates in a CT log.
type Watcher struct {
	workers    []*worker
	workersMu  sync.RWMutex
	wg         sync.WaitGroup
	context    context.Context
	certChan   chan models.Entry
	cancelFunc context.CancelFunc
}

// NewWatcher creates a new Watcher.
func NewWatcher(certChan chan models.Entry) *Watcher {
	return &Watcher{
		certChan: certChan,
	}
}

// Start starts the watcher. This method is blocking.
func (w *Watcher) Start() {
	w.context, w.cancelFunc = context.WithCancel(context.Background())

	// Create new certChan if it doesn't exist yet
	if w.certChan == nil {
		w.certChan = make(chan models.Entry, 5000)
	}

	if config.AppConfig.General.Recovery.Enabled {
		ctIndexFilePath, err := filepath.Abs(config.AppConfig.General.Recovery.CTIndexFile)
		if err != nil {
			log.Printf("Error getting absolute path for CT index file: '%s', %s\n", config.AppConfig.General.Recovery.CTIndexFile, err)
			return
		}
		// Load Saved CT Indexes
		metrics.Metrics.LoadCTIndex(ctIndexFilePath)
		// Save CTIndexes at regular intervals
		go metrics.Metrics.SaveCertIndexesAtInterval(time.Second*30, ctIndexFilePath) // save indexes every X seconds
	}

	// initialize the watcher with currently available logs
	w.updateLogs()

	log.Println("Started CT watcher")
	go certHandler(w.certChan)
	go w.watchNewLogs()

	// Wait for all workers to finish
	w.wg.Wait()
	close(w.certChan)
}

// watchNewLogs monitors the ct log list for new logs and starts a worker for each new log found.
// This method is blocking. It can be stopped by cancelling the context.
func (w *Watcher) watchNewLogs() {
	// Check for new logs once every hour
	ticker := time.NewTicker(1 * time.Hour)
	for {
		select {
		case <-ticker.C:
			w.updateLogs()
		case <-w.context.Done():
			ticker.Stop()
			return
		}
	}
}

// updateLogs checks the transparency log list for new logs and adds new workers for those to the watcher.
func (w *Watcher) updateLogs() {
	// Get a list of urls of all CT logs
	logList, err := getAllLogs()
	if err != nil {
		log.Println(err)
		return
	}

	log.Println("Checking for new ct logs...")

	// Track all URLs that should be monitored after reconciliation
	monitoredURLs := make(map[string]struct{})
	newCTs := 0

	w.workersMu.Lock()
	defer w.workersMu.Unlock()

	for _, operator := range logList.Operators {
		// Classic logs
		for _, transparencyLog := range operator.Logs {
			url := transparencyLog.URL
			desc := transparencyLog.Description
			normURL := normalizeCtlogURL(url)

			if transparencyLog.State.LogStatus() == loglist3.RetiredLogStatus {
				log.Printf("Skipping retired CT log: %s\n", normURL)
				continue
			}

			monitoredURLs[normURL] = struct{}{}
			if w.addLogIfNew(operator.Name, desc, url, false) {
				newCTs++
			}
		}

		// Tiled logs
		for _, transparencyLog := range operator.TiledLogs {
			url := transparencyLog.MonitoringURL
			desc := transparencyLog.Description
			normURL := normalizeCtlogURL(url)

			if transparencyLog.State.LogStatus() == loglist3.RetiredLogStatus {
				log.Printf("Skipping retired CT log: %s\n", normURL)
				continue
			}

			monitoredURLs[normURL] = struct{}{}
			if w.addLogIfNew(operator.Name, desc, url, true) {
				newCTs++
			}
		}
	}

	log.Printf("New ct logs found: %d\n", newCTs)

	// Optionally stop workers for logs not in the monitoredURLs set
	if *config.AppConfig.General.DropOldLogs {
		removed := 0
		for _, ctWorker := range w.workers {
			normURL := normalizeCtlogURL(ctWorker.ctURL)
			if _, ok := monitoredURLs[normURL]; !ok {
				log.Printf("Stopping worker. CT URL not found in LogList or retired: '%s'\n", ctWorker.ctURL)
				ctWorker.stop()
				removed++
			}
		}
		log.Printf("Removed ct logs: %d\n", removed)
	}

	log.Printf("Currently monitored ct logs: %d\n", len(w.workers))
}

// addLogIfNew checks if a log is already being watched and adds it if not.
// Returns true if a new log was added, false otherwise.
func (w *Watcher) addLogIfNew(operatorName, description, url string, isTiled bool) bool {
	normURL := normalizeCtlogURL(url)

	// Check if the log is already being watched
	for _, ctWorker := range w.workers {
		workerURL := normalizeCtlogURL(ctWorker.ctURL)
		if workerURL == normURL {
			return false
		}
	}

	// Log is not being watched, so add it
	w.wg.Add(1)

	lastCTIndex := metrics.Metrics.GetCTIndex(normURL)
	ctWorker := worker{
		name:         description,
		operatorName: operatorName,
		ctURL:        url,
		entryChan:    w.certChan,
		ctIndex:      lastCTIndex,
		isTiled:      isTiled,
	}
	w.workers = append(w.workers, &ctWorker)
	metrics.Metrics.Init(operatorName, normURL)

	// Start a goroutine for each worker
	go func() {
		defer w.wg.Done()
		ctWorker.startDownloadingCerts(w.context)
		w.discardWorker(&ctWorker)
	}()

	return true
}

// discardWorker removes a worker from the watcher's list of workers.
// This needs to be done when a worker stops.
func (w *Watcher) discardWorker(worker *worker) {
	log.Println("Removing worker for CT log:", worker.ctURL)

	w.workersMu.Lock()
	defer w.workersMu.Unlock()

	for i, wo := range w.workers {
		if wo == worker {
			w.workers = append(w.workers[:i], w.workers[i+1:]...)
			return
		}
	}
}

// Stop stops the watcher.
func (w *Watcher) Stop() {
	log.Printf("Stopping watcher\n")

	if config.AppConfig.General.Recovery.Enabled {
		// Store current CT Indexes before shutting down
		filePath := config.AppConfig.General.Recovery.CTIndexFile
		metrics.Metrics.SaveCertIndexes(filePath)
	}

	w.cancelFunc()
}

// CreateIndexFile creates a ct_index.json file based on the current STHs of all availble logs.
func (w *Watcher) CreateIndexFile(filePath string) error {
	logs, err := getAllLogs()
	if err != nil {
		return err
	}

	w.context, w.cancelFunc = context.WithCancel(context.Background())
	log.Println("Fetching current STH for all logs...")
	for _, operator := range logs.Operators {
		// Iterate over each log of the operator
		for _, transparencyLog := range operator.Logs {
			if transparencyLog.State.LogStatus() == loglist3.RetiredLogStatus {
				log.Printf("Skipping retired CT log: %s\n", transparencyLog.URL)
				continue
			}

			normalizedURL := normalizeCtlogURL(transparencyLog.URL)
			// Check if the log is already being watched
			metrics.Metrics.Init(operator.Name, normalizedURL)
			log.Println("Fetching STH for", normalizedURL)

			hc := http.Client{Timeout: 5 * time.Second}
			jsonClient, e := client.New(transparencyLog.URL, &hc, jsonclient.Options{UserAgent: userAgent})
			if e != nil {
				log.Printf("Error creating JSON client: %s\n", e)
				continue
			}

			sth, getSTHerr := jsonClient.GetSTH(w.context)
			if getSTHerr != nil {
				// TODO this can happen due to a 429 error. We should retry the request
				log.Printf("Could not get STH for '%s': %s\n", transparencyLog.URL, getSTHerr)
				continue
			}

			metrics.Metrics.SetCTIndex(normalizedURL, sth.TreeSize)
		}
		for _, transparencyLog := range operator.TiledLogs {
			if transparencyLog.State.LogStatus() == loglist3.RetiredLogStatus {
				log.Printf("Skipping retired CT log: %s\n", transparencyLog.MonitoringURL)
				continue
			}
			// Check if the log is already being watched
			normalizedURL := normalizeCtlogURL(transparencyLog.MonitoringURL)
			metrics.Metrics.Init(operator.Name, normalizedURL)
			log.Println("Fetching checkpoint for", normalizedURL)

			hc := &http.Client{Timeout: 10 * time.Second}
			checkpoint, fetchErr := FetchCheckpoint(w.context, hc, transparencyLog.MonitoringURL)
			if fetchErr != nil {
				log.Printf("Could not get checkpoint for '%s': %s\n", transparencyLog.MonitoringURL, fetchErr)
				return errFetchingSTHFailed
			}

			metrics.Metrics.SetCTIndex(normalizedURL, checkpoint.Size)
		}
	}
	w.cancelFunc()

	metrics.Metrics.SaveCertIndexes(filePath)
	log.Println("Index file saved to", filePath)

	return nil
}

// A worker processes a single CT log.
type worker struct {
	name         string
	operatorName string
	ctURL        string
	entryChan    chan models.Entry
	ctIndex      uint64
	mu           sync.Mutex
	running      bool
	cancel       context.CancelFunc
	isTiled      bool
}

// startDownloadingCerts starts downloading certificates from the CT log. This method is blocking.
func (w *worker) startDownloadingCerts(ctx context.Context) {
	ctx, w.cancel = context.WithCancel(ctx)

	// Normalize CT URL. We remove trailing slashes and prepend "https://" if it's not already there.
	w.ctURL = strings.TrimRight(w.ctURL, "/")
	if !strings.HasPrefix(w.ctURL, "https://") && !strings.HasPrefix(w.ctURL, "http://") {
		w.ctURL = "https://" + w.ctURL
	}

	log.Printf("Initializing worker for CT log: %s\n", w.ctURL)
	defer log.Printf("Stopping worker for CT log: %s\n", w.ctURL)

	w.mu.Lock()
	if w.running {
		log.Printf("Worker for '%s' already running\n", w.ctURL)
		w.mu.Unlock()

		return
	}

	w.running = true
	defer func() { w.running = false }()
	w.mu.Unlock()

	for {
		log.Printf("Starting worker for CT log: %s\n", w.ctURL)

		var workerErr error
		if w.isTiled {
			workerErr = w.runTiledWorker(ctx)
		} else {
			workerErr = w.runStandardWorker(ctx)
		}

		if workerErr != nil {
			if errors.Is(workerErr, errFetchingSTHFailed) {
				// TODO this could happen due to a 429 error. We should retry the request
				log.Printf("Worker for '%s' failed - could not fetch STH\n", w.ctURL)
				return
			} else if errors.Is(workerErr, errCreatingClient) {
				log.Printf("Worker for '%s' failed - could not create client\n", w.ctURL)
				return
			} else if strings.Contains(workerErr.Error(), "no such host") {
				log.Printf("Worker for '%s' failed to resolve host: %s\n", w.ctURL, workerErr)
				return
			}

			log.Printf("Worker for '%s' failed with unexpected error: %s\n", w.ctURL, workerErr)
		}

		// Check if the context was canceled
		select {
		case <-ctx.Done():
			log.Printf("Context was cancelled; Stopping worker for '%s'\n", w.ctURL)

			return
		default:
			log.Printf("Worker for '%s' sleeping for 5 seconds due to error\n", w.ctURL)
			time.Sleep(5 * time.Second)
			log.Printf("Restarting worker for '%s'\n", w.ctURL)

			continue
		}
	}
}

func (w *worker) stop() {
	w.mu.Lock()
	defer w.mu.Unlock()

	w.cancel()
}

// runStandardWorker runs the worker for a single standard CT log. This method is blocking.
func (w *worker) runStandardWorker(ctx context.Context) error {
	hc := http.Client{Timeout: 30 * time.Second}
	jsonClient, e := client.New(w.ctURL, &hc, jsonclient.Options{UserAgent: userAgent})
	if e != nil {
		log.Printf("Error creating JSON client: %s\n", e)
		return errCreatingClient
	}

	// If recovery is enabled, we start at the saved index. Otherwise, we start at the latest STH.
	recoveryEnabled := config.AppConfig.General.Recovery.Enabled
	if !recoveryEnabled {
		sth, getSTHerr := jsonClient.GetSTH(ctx)
		if getSTHerr != nil {
			// TODO this can happen due to a 429 error. We should retry the request
			log.Printf("Could not get STH for '%s': %s\n", w.ctURL, getSTHerr)
			return errFetchingSTHFailed
		}
		// Start at the latest STH to skip all the past certificates
		w.ctIndex = sth.TreeSize
	}

	certScanner := scanner.NewScanner(jsonClient, scanner.ScannerOptions{
		FetcherOptions: scanner.FetcherOptions{
			BatchSize:     100,
			ParallelFetch: 1,
			StartIndex:    int64(w.ctIndex),
			Continuous:    true,
		},
		Matcher:     scanner.MatchAll{},
		PrecertOnly: false,
		NumWorkers:  1,
		BufferSize:  config.AppConfig.General.BufferSizes.CTLog,
	})

	scanErr := certScanner.Scan(ctx, w.foundCertCallback, w.foundPrecertCallback)
	if scanErr != nil {
		log.Println("Scan error: ", scanErr)
		return scanErr
	}

	log.Printf("Exiting worker %s without error!\n", w.ctURL)

	return nil
}

// runTiledWorker runs the worker for a single tiled CT log. This method is blocking.
func (w *worker) runTiledWorker(ctx context.Context) error {
	hc := &http.Client{Timeout: 30 * time.Second}

	// If recovery is enabled and the CT index is set, we start at the saved index. Otherwise, we start at the latest checkpoint.
	validSavedCTIndexExists := config.AppConfig.General.Recovery.Enabled
	if !validSavedCTIndexExists {
		checkpoint, err := FetchCheckpoint(ctx, hc, w.ctURL)
		if err != nil {
			log.Printf("Could not get checkpoint for '%s': %s\n", w.ctURL, err)
			return errFetchingSTHFailed
		}
		// Start at the latest checkpoint to skip all the past certificates
		w.ctIndex = checkpoint.Size
	}

	// Initialize backoff for polling
	pollBackoff := &backoff.Backoff{
		Min:    1 * time.Second,
		Max:    30 * time.Second,
		Factor: 2,
		Jitter: true,
	}

	// Continuous monitoring loop
	for {
		hadNewEntries, err := w.processTiledLogUpdates(ctx, hc)
		if err != nil {
			log.Printf("Error processing tiled log updates for '%s': %s\n", w.ctURL, err)
			return err
		}

		// Reset backoff if we found new entries
		if hadNewEntries {
			pollBackoff.Reset()
		}

		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(pollBackoff.Duration()):
			// Continue to the next iteration
		}
	}
}

// processTiledLogUpdates checks for new entries in the tiled log and processes them
func (w *worker) processTiledLogUpdates(ctx context.Context, hc *http.Client) (bool, error) {
	// Fetch current checkpoint
	checkpoint, err := FetchCheckpoint(ctx, hc, w.ctURL)
	if err != nil {
		return false, fmt.Errorf("fetching checkpoint: %w", err)
	}

	currentTreeSize := checkpoint.Size
	if currentTreeSize <= w.ctIndex {
		// No new entries
		return false, nil
	}

	// Process entries from current index to new tree size
	startTile := w.ctIndex / TileSize
	endTile := currentTreeSize / TileSize

	// Process complete tiles
	for tileIndex := startTile; tileIndex < endTile; tileIndex++ {
		if err := w.processTile(ctx, hc, tileIndex, 0); err != nil {
			return false, fmt.Errorf("processing tile %d: %w", tileIndex, err)
		}
	}

	// Process partial tile if exists
	partialSize := currentTreeSize % TileSize
	if partialSize > 0 {
		if err := w.processTile(ctx, hc, endTile, partialSize); err != nil {
			log.Printf("Warning: error processing partial tile %d: %s\n", endTile, err)
			// Don't return error for partial tiles as they might be incomplete
		}
	}

	return true, nil
}

// processTile processes a single tile from the tiled log.
// partialWidth of 0 means full tile, otherwise fetch partial tile with that width.
func (w *worker) processTile(ctx context.Context, hc *http.Client, tileIndex uint64, partialWidth uint64) error {
	leaves, err := FetchTile(ctx, hc, w.ctURL, tileIndex, partialWidth)
	if err != nil {
		return fmt.Errorf("fetching tile: %w", err)
	}

	// Calculate the starting index for entries in this tile
	baseIndex := tileIndex * TileSize

	for i, leaf := range leaves {
		entryIndex := baseIndex + uint64(i)

		// Skip entries we've already processed
		if entryIndex <= w.ctIndex {
			continue
		}

		// Convert TileLeaf to RawLogEntry for compatibility with existing parsing
		rawEntry := ConvertTileLeafToRawLogEntry(leaf, entryIndex)

		// Process the entry using existing callbacks
		switch leaf.EntryType {
		case 0:
			w.foundCertCallback(rawEntry)
		case 1:
			w.foundPrecertCallback(rawEntry)
		default:
			log.Printf("Unknown entry type %d in tile %d, skipping entry at index %d\n", leaf.EntryType, tileIndex, entryIndex)
		}

		// Update the index
		w.ctIndex = entryIndex
	}

	return nil
}

// foundCertCallback is the callback that handles cases where new regular certs are found.
func (w *worker) foundCertCallback(rawEntry *ct.RawLogEntry) {
	logType := models.SourceIsRFC6962
	if w.isTiled {
		logType = models.SourceIsTiled
	}

	entry, parseErr := ParseCertstreamEntry(rawEntry, w.operatorName, w.name, w.ctURL, logType)
	if parseErr != nil {
		log.Println("Error parsing certstream entry: ", parseErr)
		return
	}

	entry.Data.UpdateType = "X509LogEntry"
	w.entryChan <- entry

	atomic.AddInt64(&metrics.ProcessedCerts, 1)
}

// foundPrecertCallback is the callback that handles cases where new precerts are found.
func (w *worker) foundPrecertCallback(rawEntry *ct.RawLogEntry) {
	logType := models.SourceIsRFC6962
	if w.isTiled {
		logType = models.SourceIsTiled
	}

	entry, parseErr := ParseCertstreamEntry(rawEntry, w.operatorName, w.name, w.ctURL, logType)
	if parseErr != nil {
		log.Println("Error parsing certstream entry: ", parseErr)
		return
	}

	entry.Data.UpdateType = "PrecertLogEntry"
	w.entryChan <- entry

	atomic.AddInt64(&metrics.ProcessedPrecerts, 1)
}

// certHandler takes the entries out of the entryChan channel and broadcasts them to all clients.
// Only a single instance of the certHandler runs per certstream server.
func certHandler(entryChan chan models.Entry) {
	var processed uint64

	for {
		entry := <-entryChan
		processed++

		if processed%1000 == 0 {
			log.Printf("Processed %d entries | Queue length: %d\n", processed, len(entryChan))
			// Every thousandth entry, we store one certificate as example
			web.SetExampleCert(entry)
		}

		// Run JSON encoding in the background and send the result to the clients.
		web.ClientHandler.Broadcast <- entry

		// Update metrics
		url := entry.Data.Source.NormalizedURL
		operator := entry.Data.Source.Operator
		index := entry.Data.CertIndex

		metrics.Metrics.Inc(operator, url, index)
	}
}

// getGoogleLogList fetches the list of all CT logs from Google Chromes CT LogList.
func getGoogleLogList() (loglist3.LogList, error) {
	// Download the list of all logs from ctLogInfo and decode JSON
	resp, err := http.Get(loglist3.LogListURL)
	if err != nil {
		return loglist3.LogList{}, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return loglist3.LogList{}, errors.New("failed to download loglist")
	}

	bodyBytes, readErr := io.ReadAll(resp.Body)
	if readErr != nil {
		log.Panic(readErr)
	}

	allLogs, parseErr := loglist3.NewFromJSON(bodyBytes)
	if parseErr != nil {
		return loglist3.LogList{}, parseErr
	}

	return *allLogs, nil
}

// getAllLogs returns a list of all CT logs.
func getAllLogs() (loglist3.LogList, error) {
	var allLogs loglist3.LogList
	var err error

	// Ability to disable default logs, if the user only wants to monitor custom logs.
	if !config.AppConfig.General.DisableDefaultLogs {
		allLogs, err = getGoogleLogList()
		if err != nil {
			log.Printf("Error fetching log list from Google: %s\n", err)
			return loglist3.LogList{}, fmt.Errorf("failed to fetch log list from Google: %w", err)
		}
	}

	// Add manually added logs from config to the allLogs list
	// if config.AppConfig.General.AdditionalLogs == nil {
	// 	return allLogs, nil
	// }

logFound:
	for _, additionalLog := range config.AppConfig.General.AdditionalLogs {
		customLog := loglist3.Log{
			URL:         additionalLog.URL,
			Description: additionalLog.Description,
		}

		operatorFound := false
		for _, operator := range allLogs.Operators {
			if operator.Name == additionalLog.Operator {
				operatorFound = true

				for _, ctlog := range operator.Logs {
					if ctlog.URL == additionalLog.URL {
						// Log already exists, skip it.
						break logFound
					}
				}
				// This works, since allLogs.Operators is a slice of pointers.
				operator.Logs = append(operator.Logs, &customLog)
				break
			}
		}

		if !operatorFound {
			newOperator := loglist3.Operator{
				Name: additionalLog.Operator,
				Logs: []*loglist3.Log{&customLog},
			}
			allLogs.Operators = append(allLogs.Operators, &newOperator)
		}
	}

	for _, additionalLog := range config.AppConfig.General.AdditionalTiledLogs {
		customLog := loglist3.TiledLog{
			MonitoringURL: additionalLog.URL,
			Description:   additionalLog.Description,
		}

		operatorFound := false

	tiledLogFound:
		for _, operator := range allLogs.Operators {
			if operator.Name == additionalLog.Operator {
				operatorFound = true
				for _, tl := range operator.TiledLogs {
					if tl.MonitoringURL == additionalLog.URL {
						// Log already exists, skip it.
						break tiledLogFound
					}
				}
				// This works, since allLogs.Operators is a slice of pointers.
				operator.TiledLogs = append(operator.TiledLogs, &customLog)
				break
			}
		}

		if !operatorFound {
			newOperator := loglist3.Operator{
				Name:      additionalLog.Operator,
				TiledLogs: []*loglist3.TiledLog{&customLog},
			}
			allLogs.Operators = append(allLogs.Operators, &newOperator)
		}
	}

	return allLogs, nil
}

func normalizeCtlogURL(input string) string {
	input = strings.TrimPrefix(input, "https://")
	input = strings.TrimPrefix(input, "http://")
	input = strings.TrimSuffix(input, "/")

	return input
}
