main.go

package main

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/url"
	"os"
	"os/signal"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"
	"tekir/crawler"
	"time"

	"github.com/fatih/color"
	flag "github.com/spf13/pflag"
)

const metadataSaveFile = "crawled_data.json"
const saveInterval = 10 * time.Second
const logFileName = "crawl.log"
const imageDir = "images"

var scoringEnabled bool // Set by --score flag to enable site scoring

// saveMetadataToFile saves the current metadata slice to the JSON file.
// Updated to accept loggers.
func saveMetadataToFile(metadata []crawler.PageMetadata, mutex *sync.Mutex, infoLogger, errorLogger *log.Logger) {
	mutex.Lock()
	// Make a copy of the metadata slice to work with after unlocking
	metadataCopy := make([]crawler.PageMetadata, len(metadata))
	copy(metadataCopy, metadata)
	mutex.Unlock() // Unlock earlier to avoid holding lock during I/O

	// Compute site scores only if enabled
	if scoringEnabled {
		metadataCopy = crawler.ScorePages(metadataCopy)
	}

	if len(metadataCopy) == 0 {
		infoLogger.Println("saveMetadataToFile: No metadata to save.")
		return
	}

	infoLogger.Printf("saveMetadataToFile: Attempting to save %d records to %s...\n", len(metadataCopy), metadataSaveFile)
	jsonData, err := json.MarshalIndent(metadataCopy, "", "  ") // Pretty print JSON
	if err != nil {
		errorLogger.Printf("Failed to marshal metadata to JSON: %v", err)
		return
	}

	// Use a temporary file for atomic write
	tempFile := metadataSaveFile + ".tmp"
	err = os.WriteFile(tempFile, jsonData, 0644)
	if err != nil {
		errorLogger.Printf("Failed to write JSON data to temporary file %s: %v", tempFile, err)
		return
	}

	// Rename the temporary file to the final destination file
	err = os.Rename(tempFile, metadataSaveFile)
	if err != nil {
		errorLogger.Printf("Failed to rename temporary file %s to %s: %v", tempFile, metadataSaveFile, err)
		// Attempt to clean up the temporary file
		os.Remove(tempFile)
		return
	}

	infoLogger.Printf("saveMetadataToFile: Successfully saved %d records to %s\n", len(metadataCopy), metadataSaveFile)
}

func main() {
	// Track start time
	startTime := time.Now()
	// Counters for pages crawled and errors
	var pagesCount int64
	var errorCount int64
	// Domain crawl counts
	var domainCount = make(map[string]int)
	var domainMutex sync.Mutex

	// --- Define and Parse Flags ---
	disobeyRobots := flag.Bool("disobey", false, "Disables robots.txt checks")
	verbose := flag.Bool("v", false, "Enable verbose logging")
	numWorkers := flag.Int("w", 5, "Number of concurrent page crawler workers")
	numImageWorkers := flag.Int("iw", 3, "Number of concurrent image download workers") // New flag
	// Score flag
	scoreFlag := flag.Bool("score", false, "Enable site scoring and include site_score in JSON output")

	// Parse command-line flags (supports flags after positional args)
	flag.Parse()
	scoringEnabled = *scoreFlag
	// Positional args (start URLs)
	args := flag.Args()

	// --- Setup Logging ---
	logFile, err := os.OpenFile(logFileName, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
	if err != nil {
		// Use standard log for this initial fatal error as loggers aren't set up yet
		log.Fatalf("Failed to open log file %s: %v", logFileName, err)
	}
	defer logFile.Close()

	// Create dedicated loggers
	// Error logger always writes to the file
	errorLogger := log.New(logFile, "ERROR: ", log.Ldate|log.Ltime|log.Lshortfile)

	// Info logger writes to file only if verbose, otherwise discards
	infoOutput := io.Discard
	if *verbose {
		infoOutput = logFile
	}
	infoLogger := log.New(infoOutput, "INFO:  ", log.Ldate|log.Ltime|log.Lshortfile)

	// If scoring-only mode is enabled, rescore existing data and exit
	if scoringEnabled {
		infoLogger.Println("Score-only mode enabled. Rescoring existing metadata...")
		// Read existing crawled data
		data, err := os.ReadFile(metadataSaveFile)
		if err != nil {
			errorLogger.Fatalf("Failed to read %s: %v", metadataSaveFile, err)
		}
		var pages []crawler.PageMetadata
		if err := json.Unmarshal(data, &pages); err != nil {
			errorLogger.Fatalf("Failed to parse JSON from %s: %v", metadataSaveFile, err)
		}
		// Recompute scores
		scored := crawler.ScorePages(pages)
		out, err := json.MarshalIndent(scored, "", "  ")
		if err != nil {
			errorLogger.Fatalf("Failed to marshal scored data: %v", err)
		}
		if err := os.WriteFile(metadataSaveFile, out, 0644); err != nil {
			errorLogger.Fatalf("Failed to write scored data to %s: %v", metadataSaveFile, err)
		}
		infoLogger.Printf("Successfully rescored %d records in %s", len(scored), metadataSaveFile)
		return
	}

	infoLogger.Println("--- Crawler Started ---")
	if *disobeyRobots {
		infoLogger.Println("Running with --disobey flag. Robots.txt checks are disabled.")
	}
	if *verbose {
		infoLogger.Println("Verbose logging enabled.")
	}
	infoLogger.Printf("Page workers: %d, Image workers: %d", *numWorkers, *numImageWorkers)

	// --- Get Start URLs ---
	if len(args) < 1 {
		fmt.Println("Usage: go run main.go [-v] [--disobey] [-w N] [-iw N] <start-url> [additional-urls...]")
		errorLogger.Fatalln("No start URL(s) provided.")
	}
	// Helper to add default scheme
	ensureScheme := func(raw string) string {
		if strings.HasPrefix(raw, "http://") || strings.HasPrefix(raw, "https://") {
			return raw
		}
		return "https://" + raw
	}
	// Normalize all provided start URLs
	startURLs := make([]string, 0, len(args))
	for _, raw := range args {
		u := ensureScheme(raw)
		startURLs = append(startURLs, u)
	}
	infoLogger.Printf("Starting crawl with URLs: %v", startURLs)

	// Use the constructor from the crawler package
	visited := crawler.NewVisited()
	linkQueue := make(chan []string)                               // Channel to receive new links found by crawlers
	metadataQueue := make(chan crawler.PageMetadata)               // Use crawler.PageMetadata
	worklist := make(chan string, *numWorkers*2)                   // Add some buffer to worklist
	imageQueue := make(chan crawler.ImageTask, *numImageWorkers*5) // Add buffer to image queue
	backlinkStore := make(map[string]map[string]bool)              // Added: map[targetURL]map[sourceURL]bool
	var backlinkMutex sync.RWMutex                                 // Added: Mutex for backlinkStore

	var wg sync.WaitGroup         // WaitGroup to wait for crawlers to finish
	var imageWg sync.WaitGroup    // WaitGroup for image download workers
	var metadataWg sync.WaitGroup // WaitGroup for metadata handler

	// Slice to store collected metadata
	var collectedMetadata []crawler.PageMetadata
	var metadataMutex sync.Mutex // Mutex to protect access to collectedMetadata

	// Load existing metadata if file exists
	if _, err := os.Stat(metadataSaveFile); err == nil {
		data, err := os.ReadFile(metadataSaveFile)
		if err != nil {
			errorLogger.Printf("Failed to read existing metadata file: %v", err)
		} else {
			var existing []crawler.PageMetadata
			err = json.Unmarshal(data, &existing)
			if err != nil {
				errorLogger.Printf("Failed to parse existing metadata JSON: %v", err)
			} else {
				collectedMetadata = existing
				infoLogger.Printf("Loaded %d existing metadata records", len(collectedMetadata))
			}
		}
	}
	// Build map of existing addresses for duplication
	existingPages := make(map[string]bool)
	for _, m := range collectedMetadata {
		existingPages[m.SiteAddress] = true
	}

	// Load or initialize persistent visited sites file
	const persistentVisitedFile = "visited_sites.json"
	var persistentVisited map[string]int64
	var persistentVisitedMutex sync.Mutex
	if _, err := os.Stat(persistentVisitedFile); err == nil {
		if data, err := os.ReadFile(persistentVisitedFile); err == nil {
			json.Unmarshal(data, &persistentVisited)
			infoLogger.Printf("Loaded %d visited sites from %s", len(persistentVisited), persistentVisitedFile)
		}
	} else {
		persistentVisited = make(map[string]int64)
	}
	// Function to atomically save persistent visited sites
	saveVisited := func() {
		persistentVisitedMutex.Lock()
		defer persistentVisitedMutex.Unlock()
		if buf, err := json.MarshalIndent(persistentVisited, "", "  "); err != nil {
			errorLogger.Printf("Failed to marshal visited sites JSON: %v", err)
		} else if tmp := persistentVisitedFile + ".tmp"; os.WriteFile(tmp, buf, 0644) == nil {
			if err := os.Rename(tmp, persistentVisitedFile); err != nil {
				errorLogger.Printf("Failed to rename visited temp file: %v", err)
			}
		} else {
			errorLogger.Printf("Failed to write temp visited file")
		}
	}

	// Setup OS signal handler for graceful shutdown
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
	go func() {
		<-sigChan
		infoLogger.Println("Received interrupt signal, saving metadata before exit...")
		saveMetadataToFile(collectedMetadata, &metadataMutex, infoLogger, errorLogger)
		saveVisited()
		os.Exit(0)
	}()

	// Ticker for periodic saving
	saveTicker := time.NewTicker(saveInterval)
	defer saveTicker.Stop() // Ensure ticker is stopped

	// Reporter function for status (UI-like panel)
	reportStatus := func() {
		// Clear screen
		fmt.Print("\033[2J\033[H")
		// Header
		color.New(color.FgCyan, color.Bold).Println("=== TekirCrawler Status ===")
		// Uptime and time
		color.New(color.FgYellow).Printf("Time: %s | Uptime: %s\n", time.Now().Format(time.RFC3339), time.Since(startTime).Truncate(time.Second))
		// Pages and errors
		color.New(color.FgGreen).Printf("Pages Crawled: %d", atomic.LoadInt64(&pagesCount))
		color.New(color.FgRed).Printf("  Errors: %d\n", atomic.LoadInt64(&errorCount))
		// Queue lengths
		color.New(color.FgMagenta).Printf("Worklist: %d   Images: %d\n", len(worklist), len(imageQueue))
		// Worker counts
		color.New(color.FgBlue).Printf("Crawler Workers: %d   Image Workers: %d\n", *numWorkers, *numImageWorkers)
		// Most crawled domain
		domainMutex.Lock()
		topDomain, topCount := "", 0
		for d, c := range domainCount {
			if c > topCount {
				topDomain, topCount = d, c
			}
		}
		domainMutex.Unlock()
		if topCount > 0 {
			color.New(color.FgGreen).Printf("Top Domain: %s (%d pages)\n", topDomain, topCount)
		}
		// Footer
		color.New(color.FgCyan, color.Bold).Println("===========================")
	}

	// Initial status print
	reportStatus()

	// Start periodic status reporting every 1 second
	statusTicker := time.NewTicker(1 * time.Second)
	defer statusTicker.Stop()
	go func() {
		for range statusTicker.C {
			reportStatus()
		}
	}()

	// Goroutine to handle metadata and periodic saving
	metadataWg.Add(1)
	go func() {
		defer metadataWg.Done()
		infoLogger.Println("Metadata handling goroutine started.")
		for {
			select {
			case meta, ok := <-metadataQueue:
				if !ok {
					infoLogger.Println("Metadata channel closed. Performing final save...")
					saveMetadataToFile(collectedMetadata, &metadataMutex, infoLogger, errorLogger)
					saveVisited()
					infoLogger.Println("Metadata goroutine finished.")
					return
				}
				if meta.Crawled > 0 {
					// populate backlinks
					backlinkMutex.RLock()
					if sources, found := backlinkStore[meta.SiteAddress]; found {
						meta.Backlinks = make([]string, 0, len(sources))
						for sourceURL := range sources {
							meta.Backlinks = append(meta.Backlinks, sourceURL)
						}
						meta.BacklinkCount = len(meta.Backlinks)
					} else {
						meta.BacklinkCount = 0
					}
					backlinkMutex.RUnlock()

					// simplified log: remove SiteScore
					crawlTime := time.Unix(meta.Crawled, 0)
					infoLogger.Printf("Metadata Processed: URL=%s, Title='%s', Desc='%s', Tags=%v, Images=%d, Favicon=%t, Backlinks=%d, CrawledAt=%s",
						meta.SiteAddress, meta.Title, meta.Description, meta.Tags, len(meta.Images), meta.FaviconPath != "", meta.BacklinkCount, crawlTime.Format(time.RFC3339))
					metadataMutex.Lock()
					// If page was loaded from existing file, append twice on first encounter
					if existingPages[meta.SiteAddress] {
						collectedMetadata = append(collectedMetadata, meta, meta)
						existingPages[meta.SiteAddress] = false
					} else {
						collectedMetadata = append(collectedMetadata, meta)
					}
					metadataMutex.Unlock()

					// record to persistent visited file
					persistentVisitedMutex.Lock()
					persistentVisited[meta.SiteAddress] = meta.Crawled
					persistentVisitedMutex.Unlock()
					saveVisited()

					// Increment counters
					atomic.AddInt64(&pagesCount, 1)
					parsedURL, err := url.Parse(meta.SiteAddress)
					if err == nil {
						domainMutex.Lock()
						domainCount[parsedURL.Host]++
						domainMutex.Unlock()
					}
				} else {
					reason := meta.Description
					if reason == "" {
						reason = "Unknown error during crawl or parse"
					}
					errorLogger.Printf("Metadata Failure: URL=%s, Reason: %s", meta.SiteAddress, reason)
					atomic.AddInt64(&errorCount, 1)
				}
			case <-saveTicker.C:
				infoLogger.Println("Periodic save triggered.")
				saveMetadataToFile(collectedMetadata, &metadataMutex, infoLogger, errorLogger)
				saveVisited()
			}
		}
	}()

	// Start Image Worker Goroutines
	infoLogger.Printf("Starting %d image workers...", *numImageWorkers)
	for i := 0; i < *numImageWorkers; i++ {
		imageWg.Add(1)
		go func(workerID int) {
			defer imageWg.Done()
			infoLogger.Printf("Image Worker %d started", workerID)
			for task := range imageQueue {
				// Use the Verbose flag from the task
				_, err := crawler.DownloadFile(task.ImageURL, imageDir, task.BaseFilename, true, task.Verbose, infoLogger, errorLogger)
				if err != nil {
					errorLogger.Printf("Image Worker %d: Failed %s: %v", workerID, task.ImageURL, err)
				}
			}
			infoLogger.Printf("Image Worker %d finished", workerID)
		}(i)
	}

	// Start worker goroutines
	infoLogger.Printf("Starting %d crawler workers...", *numWorkers)
	for i := 0; i < *numWorkers; i++ {
		wg.Add(1)
		go func(workerID int) {
			defer wg.Done()
			infoLogger.Printf("Worker %d started", workerID)
			for urlToCrawl := range worklist {
				infoLogger.Printf("Worker %d: Processing %s", workerID, urlToCrawl)
				// Pass backlinkStore and mutex to Crawl
				crawler.Crawl(urlToCrawl, visited, linkQueue, metadataQueue, imageQueue, backlinkStore, &backlinkMutex, *disobeyRobots, *verbose, infoLogger, errorLogger)
			}
			infoLogger.Printf("Worker %d finished", workerID)
		}(i)
	}

	// Seed worklist with all start URLs
	seen := make(map[string]bool)
	activeCrawlers := 0
	for _, u := range startURLs {
		worklist <- u
		seen[u] = true
		activeCrawlers++
		infoLogger.Printf("Queueing start URL: %s", u)
	}

	// Keep track of URLs being processed or waiting
	for activeCrawlers > 0 {
		links := <-linkQueue
		activeCrawlers-- // Decrement when a crawl task *finishes* (provides links)
		infoLogger.Printf("Crawler task finished, %d active crawlers potentially remaining", activeCrawlers)
		for _, link := range links {
			if !seen[link] {
				// skip if already in persistent visited file
				persistentVisitedMutex.Lock()
				visitedBefore := persistentVisited[link] // visitedBefore is the timestamp (int64) or 0 if not found
				persistentVisitedMutex.Unlock()
				if visitedBefore == 0 && len(link) > 0 && strings.HasPrefix(link, "http") {
					seen[link] = true
					activeCrawlers++
					infoLogger.Printf("Queueing new link: %s (%d active crawlers)", link, activeCrawlers)
					go func(l string) { worklist <- l }(link)
				} else {
					infoLogger.Printf("Skipping already visited link: %s", link)
				}
			}
		}
	}

	infoLogger.Println("Crawling loop finished. Closing worklist and waiting for workers...")

	close(worklist)
	wg.Wait()
	infoLogger.Println("All workers finished. Closing image queue and waiting for image workers...")
	close(imageQueue)
	imageWg.Wait()
	infoLogger.Println("All image workers finished. Closing metadata queue and waiting for metadata handler...")
	close(metadataQueue)
	infoLogger.Println("Waiting for metadata handler to finish...")
	metadataWg.Wait()

	infoLogger.Println("Crawling finished.")

	// --- Final Summary (keep on console) ---
	fmt.Println("--- Crawl Summary ---")
	fmt.Printf("Visited %d pages.\n", visited.Count())

	metadataMutex.Lock()
	finalCount := len(collectedMetadata)
	metadataMutex.Unlock()

	if finalCount > 0 {
		fmt.Printf("Finished with %d total metadata records saved to %s.\n", finalCount, metadataSaveFile)
	} else {
		fmt.Println("No successful metadata was collected.")
	}
	fmt.Printf("Detailed logs written to %s (Errors always logged, Info logged if -v enabled)\n", logFileName)
	infoLogger.Println("--- Crawler Finished --- ")
	errorLogger.Println("--- Crawler Finished (Error Log Marker) --- ")
}