package main

import (
	"crypto/tls"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"path/filepath"
	"regexp"
	"strings"
	"sync"
	"syscall"
	"time"

	"golang.org/x/net/html"
)

// Configuration for the crawler
type Config struct {
	BaseURL         string
	OutputDir       string
	Concurrency     int
	Timeout         time.Duration
	AllowedDomains  []string
	URLReplaceRules []ReplaceRule
	SkipTLS         bool
	UserAgent       string
	QueueFile       string        // 新增：保存队列的文件路径
	SaveInterval    time.Duration // 新增：定期保存队列的时间间隔
}

// Rule for replacing URLs
type ReplaceRule struct {
	Pattern     *regexp.Regexp
	Replacement string
}

// CrawlState 存储爬虫状态
type CrawlState struct {
	Queue       []string        `json:"queue"`
	VisitedURLs map[string]bool `json:"visited_urls"`
	BaseURL     string          `json:"base_url"`
	BaseURLPath string          `json:"base_url_path"`
	Timestamp   time.Time       `json:"timestamp"`
}

// Crawler struct
type Crawler struct {
	config      Config
	visitedURLs map[string]bool
	queue       []string
	client      *http.Client
	mutex       sync.Mutex
	wg          sync.WaitGroup
	baseURLPath string    // 存储基础URL的路径，用于防止爬取父路径
	shutdown    chan bool // 新增：用于通知goroutine关闭的通道
	stopped     bool      // 新增：标记爬虫是否已停止
}

// Create a new crawler
func NewCrawler(config Config) *Crawler {
	// Initialize HTTP client with timeout and TLS configuration
	transport := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: config.SkipTLS},
	}

	client := &http.Client{
		Timeout:   config.Timeout,
		Transport: transport,
	}

	// 解析基础URL的路径，用于后续比较
	baseURL, _ := url.Parse(config.BaseURL)
	basePath := baseURL.Path
	if !strings.HasSuffix(basePath, "/") {
		// 确保路径以"/"结尾，方便后续处理
		basePath = basePath + "/"
	}

	return &Crawler{
		config:      config,
		visitedURLs: make(map[string]bool),
		client:      client,
		baseURLPath: basePath,
		shutdown:    make(chan bool),
		stopped:     false,
	}
}

// 新增：从文件加载之前的爬虫状态
func (c *Crawler) LoadState() error {
	if c.config.QueueFile == "" {
		return nil // 如果没有指定队列文件，则不加载
	}

	// 检查文件是否存在
	if _, err := os.Stat(c.config.QueueFile); os.IsNotExist(err) {
		fmt.Println("Queue file does not exist, starting fresh crawl")
		return nil
	}

	// 读取文件内容
	data, err := ioutil.ReadFile(c.config.QueueFile)
	if err != nil {
		return fmt.Errorf("failed to read queue file: %v", err)
	}

	// 解析JSON
	var state CrawlState
	err = json.Unmarshal(data, &state)
	if err != nil {
		return fmt.Errorf("failed to parse queue file: %v", err)
	}

	// 验证基础URL是否一致
	if state.BaseURL != c.config.BaseURL {
		fmt.Printf("Warning: Base URL in queue file (%s) differs from current base URL (%s)\n",
			state.BaseURL, c.config.BaseURL)
		// 可以选择继续或返回错误
	}

	// 恢复状态
	c.mutex.Lock()
	c.queue = state.Queue
	c.visitedURLs = state.VisitedURLs
	c.mutex.Unlock()

	fmt.Printf("Loaded %d URLs in queue and %d visited URLs from %s\n",
		len(state.Queue), len(state.VisitedURLs), c.config.QueueFile)

	return nil
}

// 新增：保存爬虫状态到文件
func (c *Crawler) SaveState() error {
	if c.config.QueueFile == "" {
		return nil // 如果没有指定队列文件，则不保存
	}

	c.mutex.Lock()
	state := CrawlState{
		Queue:       c.queue,
		VisitedURLs: c.visitedURLs,
		BaseURL:     c.config.BaseURL,
		BaseURLPath: c.baseURLPath,
		Timestamp:   time.Now(),
	}
	c.mutex.Unlock()

	// 将状态转换为JSON
	data, err := json.MarshalIndent(state, "", "  ")
	if err != nil {
		return fmt.Errorf("failed to marshal state: %v", err)
	}

	// 确保目录存在
	dir := filepath.Dir(c.config.QueueFile)
	if dir != "" && dir != "." {
		if err := os.MkdirAll(dir, 0755); err != nil {
			return fmt.Errorf("failed to create directory for queue file: %v", err)
		}
	}

	// 写入文件
	if err := ioutil.WriteFile(c.config.QueueFile, data, 0644); err != nil {
		return fmt.Errorf("failed to write queue file: %v", err)
	}

	fmt.Printf("Saved %d URLs in queue and %d visited URLs to %s\n",
		len(state.Queue), len(state.VisitedURLs), c.config.QueueFile)

	return nil
}

// 新增：定期保存状态的goroutine
func (c *Crawler) autoSaveWorker() {
	ticker := time.NewTicker(c.config.SaveInterval)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			// 定期保存状态
			if err := c.SaveState(); err != nil {
				fmt.Printf("Error auto-saving state: %v\n", err)
			}
		case <-c.shutdown:
			// 收到关闭信号
			return
		}
	}
}

// Start crawling process
func (c *Crawler) Start() error {
	// 尝试加载之前的状态
	if err := c.LoadState(); err != nil {
		fmt.Printf("Warning: Failed to load previous state: %v\n", err)
	}

	baseURL, err := url.Parse(c.config.BaseURL)
	if err != nil {
		return fmt.Errorf("invalid base URL: %v", err)
	}

	// Create output directory
	err = os.MkdirAll(c.config.OutputDir, 0755)
	if err != nil {
		return fmt.Errorf("failed to create output directory: %v", err)
	}

	// 如果队列为空，添加基础URL
	c.mutex.Lock()
	if len(c.queue) == 0 {
		c.queue = append(c.queue, baseURL.String())
		c.visitedURLs[baseURL.String()] = true
	}
	c.mutex.Unlock()

	// 启动自动保存goroutine
	if c.config.SaveInterval > 0 {
		go c.autoSaveWorker()
	}

	// 设置优雅关闭的处理
	setupSignalHandler(c)

	// Create worker goroutines
	for i := 0; i < c.config.Concurrency; i++ {
		c.wg.Add(1)
		go c.worker()
	}

	c.wg.Wait()

	// 停止自动保存goroutine
	if c.config.SaveInterval > 0 {
		c.shutdown <- true
	}

	// 保存最终状态
	c.stopped = true
	return c.SaveState()
}

// Worker goroutine
func (c *Crawler) worker() {
	defer c.wg.Done()

	for {
		// 检查是否应该停止
		c.mutex.Lock()
		if c.stopped {
			c.mutex.Unlock()
			return
		}
		c.mutex.Unlock()

		// Get next URL to process
		targetURL := c.getNextURL()
		if targetURL == "" {
			return
		}

		// Process the URL
		c.processURL(targetURL)
	}
}

// Get next URL from queue (thread-safe)
func (c *Crawler) getNextURL() string {
	c.mutex.Lock()
	defer c.mutex.Unlock()

	if len(c.queue) == 0 {
		return ""
	}

	// Get next URL and remove it from queue
	nextURL := c.queue[0]
	c.queue = c.queue[1:]
	return nextURL
}

// 检查URL是否是父路径
func (c *Crawler) isParentPath(targetURL string) bool {
	parsedURL, err := url.Parse(targetURL)
	if err != nil {
		return false
	}

	// 获取目标URL的路径
	urlPath := parsedURL.Path
	if !strings.HasSuffix(urlPath, "/") {
		urlPath = urlPath + "/"
	}

	// 检查目标URL的路径是否是基础URL路径的父路径
	baseURL, _ := url.Parse(c.config.BaseURL)
	if parsedURL.Host == baseURL.Host && len(urlPath) < len(c.baseURLPath) {
		// 如果在同一域名下，且路径长度更短，可能是父路径
		return true
	}

	return false
}

// Add URL to queue (thread-safe)
func (c *Crawler) addURL(targetURL string) {
	parsedURL, err := url.Parse(targetURL)
	if err != nil {
		return
	}

	// Make sure URL is absolute
	baseURL, _ := url.Parse(c.config.BaseURL)
	absoluteURL := baseURL.ResolveReference(parsedURL).String()

	// Clean the URL - remove fragments
	cleanURL, err := url.Parse(absoluteURL)
	if err != nil {
		return
	}
	cleanURL.Fragment = ""
	absoluteURL = cleanURL.String()

	// Skip if URL is already visited
	c.mutex.Lock()
	defer c.mutex.Unlock()

	if c.visitedURLs[absoluteURL] {
		return
	}

	// Check if URL is in allowed domains
	if !c.isAllowedDomain(absoluteURL) {
		return
	}

	// Mark URL as visited and add to queue
	c.visitedURLs[absoluteURL] = true
	c.queue = append(c.queue, absoluteURL)
	fmt.Printf("Added to queue: %s\n", absoluteURL)
}

// Check if URL is in allowed domains
func (c *Crawler) isAllowedDomain(targetURL string) bool {
	parsedURL, err := url.Parse(targetURL)
	if err != nil {
		return false
	}

	// If no domains specified, only allow the base domain
	if len(c.config.AllowedDomains) == 0 {
		baseURL, _ := url.Parse(c.config.BaseURL)
		return parsedURL.Host == baseURL.Host
	}

	// Check against allowed domains
	for _, domain := range c.config.AllowedDomains {
		if parsedURL.Host == domain {
			return true
		}
	}

	return false
}

// 新增：设置关闭爬虫的方法
func (c *Crawler) Stop() {
	c.mutex.Lock()
	c.stopped = true
	c.mutex.Unlock()

	// 保存状态
	if err := c.SaveState(); err != nil {
		fmt.Printf("Error saving state during shutdown: %v\n", err)
	}
}

// Process a URL
func (c *Crawler) processURL(targetURL string) {
    // 因为一些原因，这里暂不公开该函数，请尝试自行完成，如果无法完成，请放弃尝试
}

// Parse links from HTML response
func (c *Crawler) parseLinksFromResponse(sourceURL, filePath string) {
	// Read the saved file
	file, err := os.Open(filePath)
	if err != nil {
		fmt.Printf("Error opening file %s: %v\n", filePath, err)
		return
	}
	defer file.Close()

	// Parse HTML
	doc, err := html.Parse(file)
	if err != nil {
		fmt.Printf("Error parsing HTML from %s: %v\n", filePath, err)
		return
	}

	// Parse base URL of the document
	baseURL, _ := url.Parse(sourceURL)

	// 检查HTML中是否有base标签
	baseHref := findBaseHref(doc)
	if baseHref != "" {
		// 如果有base标签，更新baseURL
		parsedBase, err := url.Parse(baseHref)
		if err == nil {
			baseURL = baseURL.ResolveReference(parsedBase)
		}
	}

	// Extract links
	var links []string
	extractLinks(doc, &links)

	// Process links
	for _, link := range links {
		// Skip non-HTTP/HTTPS links
		if strings.HasPrefix(link, "mailto:") || strings.HasPrefix(link, "tel:") ||
			strings.HasPrefix(link, "javascript:") || strings.HasPrefix(link, "#") {
			continue
		}

		// Resolve URL
		parsedURL, err := url.Parse(link)
		if err != nil {
			continue
		}

		absoluteURL := baseURL.ResolveReference(parsedURL).String()
		c.addURL(absoluteURL)
	}
}

// 查找HTML文档中的base标签
func findBaseHref(n *html.Node) string {
	if n.Type == html.ElementNode && n.Data == "base" {
		for _, attr := range n.Attr {
			if attr.Key == "href" {
				return attr.Val
			}
		}
	}

	// 递归处理子节点
	for c := n.FirstChild; c != nil; c = c.NextSibling {
		if href := findBaseHref(c); href != "" {
			return href
		}
	}

	return ""
}

// Extract links from HTML node
func extractLinks(n *html.Node, links *[]string) {
	if n.Type == html.ElementNode {
		var attrName string

		// Determine which attribute to look for based on the element type
		switch n.Data {
		case "a", "link":
			attrName = "href"
		case "img", "script", "iframe", "source", "embed":
			attrName = "src"
		case "form":
			attrName = "action"
		case "object":
			attrName = "data"
		case "area":
			attrName = "href"
		}

		// Extract attribute value if it exists
		if attrName != "" {
			for _, attr := range n.Attr {
				if attr.Key == attrName {
					*links = append(*links, attr.Val)
					break
				}
			}
		}
	}

	// Recursively process child nodes
	for c := n.FirstChild; c != nil; c = c.NextSibling {
		extractLinks(c, links)
	}
}

// 新增：捕获操作系统信号，用于优雅关闭
func setupSignalHandler(c *Crawler) {
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		<-sigChan
		fmt.Println("\nReceived shutdown signal. Saving state and exiting...")
		c.Stop()
	}()
}

func main() {
	// Parse command line flags
	baseURLFlag := flag.String("url", "", "Base URL to crawl (required)")
	outputDirFlag := flag.String("output", "downloaded_site", "Output directory")
	concurrencyFlag := flag.Int("concurrency", 5, "Number of concurrent workers")
	timeoutFlag := flag.Int("timeout", 30, "Request timeout in seconds")
	skipTLSFlag := flag.Bool("skip-tls", false, "Skip TLS certificate verification")
	userAgentFlag := flag.String("user-agent", "Mozilla/5.0 (compatible; GoCrawler/1.0)", "User-Agent string")

	// 新增命令行参数
	queueFileFlag := flag.String("queue-file", "", "File to save/load queue state")
	saveIntervalFlag := flag.Int("save-interval", 300, "Interval (in seconds) to auto-save queue state (0 to disable)")

	// URL replacement rules
	replaceFlag := flag.String("replace", "", "URL replacement rules (pattern:replacement,pattern:replacement,...)")

	// Allowed domains
	domainsFlag := flag.String("domains", "", "Allowed domains to crawl (comma-separated)")

	flag.Parse()

	// Validate base URL
	if *baseURLFlag == "" {
		fmt.Println("Error: Base URL is required")
		flag.Usage()
		os.Exit(1)
	}

	// Parse allowed domains
	var allowedDomains []string
	if *domainsFlag != "" {
		allowedDomains = strings.Split(*domainsFlag, ",")
	}

	// Parse replacement rules
	var replaceRules []ReplaceRule
	if *replaceFlag != "" {
		rules := strings.Split(*replaceFlag, ",")
		for _, rule := range rules {
			parts := strings.SplitN(rule, ":", 2)
			if len(parts) != 2 {
				fmt.Printf("Warning: Invalid replacement rule: %s\n", rule)
				continue
			}

			pattern, err := regexp.Compile(parts[0])
			if err != nil {
				fmt.Printf("Warning: Invalid regex pattern: %s\n", parts[0])
				continue
			}

			replaceRules = append(replaceRules, ReplaceRule{
				Pattern:     pattern,
				Replacement: parts[1],
			})
		}
	}

	// 默认队列文件路径（如果未指定）
	queueFile := *queueFileFlag
	if queueFile == "" && *baseURLFlag != "" {
		// 从 baseURL 生成默认的队列文件名
		baseURL, err := url.Parse(*baseURLFlag)
		if err == nil {
			queueFile = filepath.Join(*outputDirFlag, fmt.Sprintf("queue_%s.json", baseURL.Hostname()))
		}
	}

	// Create configuration
	config := Config{
		BaseURL:         *baseURLFlag,
		OutputDir:       *outputDirFlag,
		Concurrency:     *concurrencyFlag,
		Timeout:         time.Duration(*timeoutFlag) * time.Second,
		AllowedDomains:  allowedDomains,
		URLReplaceRules: replaceRules,
		SkipTLS:         *skipTLSFlag,
		UserAgent:       *userAgentFlag,
		QueueFile:       queueFile,                                      // 新增：队列文件路径
		SaveInterval:    time.Duration(*saveIntervalFlag) * time.Second, // 新增：保存间隔
	}

	// Create and start crawler
	crawler := NewCrawler(config)
	err := crawler.Start()
	if err != nil {
		fmt.Printf("Error: %v\n", err)
		os.Exit(1)
	}

	fmt.Println("Crawl completed successfully")
}