package main

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

	_ "github.com/go-sql-driver/mysql"
	"golang.org/x/net/publicsuffix"
)

func main() {
	// Connect to database
	db, err := sql.Open("mysql", "sp:0000@tcp(l2:3306)/sp_spider")
	if err != nil {
		log.Fatalf("Failed to connect to database: %v", err)
	}
	defer db.Close()

	// Test connection
	if err := db.Ping(); err != nil {
		log.Fatalf("Failed to ping database: %v", err)
	}

	// Query only unprocessed domains from unique_domains table
	rows, err := db.Query("SELECT id, domain FROM unique_domains WHERE valid IS NULL")
	if err != nil {
		log.Fatalf("Failed to query domains: %v", err)
	}
	defer rows.Close()

	// Process each domain one by one
	for rows.Next() {
		var id int
		var domain string
		if err := rows.Scan(&id, &domain); err != nil {
			log.Printf("Failed to scan domain: %v", err)
			continue
		}

		// Process domain
		domainType := processDomain(domain)

		// Update the domain record
		updateQuery := "UPDATE unique_domains SET valid = ? WHERE id = ?"
		_, err = db.Exec(updateQuery, domainType, id)
		if err != nil {
			log.Printf("Failed to update domain %s: %v", domain, err)
		} else {
			fmt.Printf("Updated ID %d: %s -> %s\n", id, domain, domainType)
		}
	}

	if err = rows.Err(); err != nil {
		log.Fatalf("Row iteration error: %v", err)
	}
}

func processDomain(domain string) string {
	domain = strings.TrimSpace(domain)
	if domain == "" {
		return "invalid"
	}

	// Get the effective TLD+1 (registrable domain)
	etld1, err := publicsuffix.EffectiveTLDPlusOne(domain)
	if err != nil {
		return "invalid"
	}

	// Determine domain type
	if domain == etld1 {
		return "valid"
	} else {
		return "sub"
	}
}