package dns

import (
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"path/filepath"
	"sort"
	"sync"
	"time"

	"dns-server/config"
	"dns-server/logger"
	"dns-server/shield"

	"github.com/miekg/dns"
)

// BlockedDomain 屏蔽域名统计

type BlockedDomain struct {
	Domain   string
	Count    int64
	LastSeen time.Time
}

// StatsData 用于持久化的统计数据结构
type StatsData struct {
	Stats              *Stats                `json:"stats"`
	BlockedDomains     map[string]*BlockedDomain `json:"blockedDomains"`
	ResolvedDomains    map[string]*BlockedDomain `json:"resolvedDomains"`
	HourlyStats        map[string]int64      `json:"hourlyStats"`
	DailyStats         map[string]int64      `json:"dailyStats"`
	MonthlyStats       map[string]int64      `json:"monthlyStats"`
	LastSaved          time.Time             `json:"lastSaved"`
}

// Server DNS服务器
type Server struct {
	config               *config.DNSConfig
	shieldConfig         *config.ShieldConfig
	shieldManager        *shield.ShieldManager
	server               *dns.Server
	resolver             *dns.Client
	ctx                  context.Context
	cancel               context.CancelFunc
	statsMutex           sync.Mutex
	stats                *Stats
	blockedDomainsMutex  sync.RWMutex
	blockedDomains       map[string]*BlockedDomain
	resolvedDomainsMutex sync.RWMutex
	resolvedDomains      map[string]*BlockedDomain // 用于记录解析的域名
	hourlyStatsMutex     sync.RWMutex
	hourlyStats          map[string]int64 // 按小时统计屏蔽数量
	dailyStatsMutex      sync.RWMutex
	dailyStats           map[string]int64 // 按天统计屏蔽数量
	monthlyStatsMutex    sync.RWMutex
	monthlyStats         map[string]int64 // 按月统计屏蔽数量
	saveTicker           *time.Ticker     // 用于定时保存数据
	saveDone             chan struct{}    // 用于通知保存协程停止
}

// Stats DNS服务器统计信息
type Stats struct {
	Queries   int64
	Blocked   int64
	Allowed   int64
	Errors    int64
	LastQuery time.Time
}

// NewServer 创建DNS服务器实例
func NewServer(config *config.DNSConfig, shieldConfig *config.ShieldConfig, shieldManager *shield.ShieldManager) *Server {
	ctx, cancel := context.WithCancel(context.Background())
	server := &Server{
		config:        config,
		shieldConfig:  shieldConfig,
		shieldManager: shieldManager,
		resolver: &dns.Client{
			Net:     "udp",
			Timeout: time.Duration(config.Timeout) * time.Millisecond,
		},
		ctx:    ctx,
		cancel: cancel,
		stats: &Stats{
			Queries: 0,
			Blocked: 0,
			Allowed: 0,
			Errors:  0,
		},
		blockedDomains:  make(map[string]*BlockedDomain),
		resolvedDomains: make(map[string]*BlockedDomain),
		hourlyStats:     make(map[string]int64),
		dailyStats:      make(map[string]int64),
		monthlyStats:    make(map[string]int64),
		saveDone:        make(chan struct{}),
	}
	
	// 加载已保存的统计数据
	server.loadStatsData()
	
	return server
	
}

// Start 启动DNS服务器
func (s *Server) Start() error {
	s.server = &dns.Server{
		Addr:    fmt.Sprintf(":%d", s.config.Port),
		Net:     "udp",
		Handler: dns.HandlerFunc(s.handleDNSRequest),
	}

	// 启动TCP服务器（用于大型响应）
	tcpServer := &dns.Server{
		Addr:    fmt.Sprintf(":%d", s.config.Port),
		Net:     "tcp",
		Handler: dns.HandlerFunc(s.handleDNSRequest),
	}

	// 启动UDP服务
	go func() {
		logger.Info(fmt.Sprintf("DNS UDP服务器启动，监听端口: %d", s.config.Port))
		if err := s.server.ListenAndServe(); err != nil {
			logger.Error("DNS UDP服务器启动失败", "error", err)
			s.cancel()
		}
	}()

	// 启动TCP服务
	go func() {
		logger.Info(fmt.Sprintf("DNS TCP服务器启动，监听端口: %d", s.config.Port))
		if err := tcpServer.ListenAndServe(); err != nil {
			logger.Error("DNS TCP服务器启动失败", "error", err)
			s.cancel()
		}
	}()

	// 等待停止信号
	<-s.ctx.Done()
	return nil
}

// Stop 停止DNS服务器
func (s *Server) Stop() {
	// 发送停止信号给保存协程
	close(s.saveDone)
	
	// 最后保存一次数据
	s.saveStatsData()
	
	// 停止服务器
	s.cancel()
	if s.server != nil {
		s.server.Shutdown()
	}
	logger.Info("DNS服务器已停止")
}

// handleDNSRequest 处理DNS请求
func (s *Server) handleDNSRequest(w dns.ResponseWriter, r *dns.Msg) {
	s.updateStats(func(stats *Stats) {
		stats.Queries++
		stats.LastQuery = time.Now()
	})

	// 只处理递归查询
	if r.RecursionDesired == false {
		response := new(dns.Msg)
		response.SetReply(r)
		response.RecursionAvailable = true
		response.SetRcode(r, dns.RcodeRefused)
		w.WriteMsg(response)
		return
	}

	// 获取查询域名
	var domain string
	if len(r.Question) > 0 {
		domain = r.Question[0].Name
		// 移除末尾的点
		if len(domain) > 0 && domain[len(domain)-1] == '.' {
			domain = domain[:len(domain)-1]
		}
	}

	logger.Debug("接收到DNS查询", "domain", domain, "type", r.Question[0].Qtype, "client", w.RemoteAddr())

	// 检查hosts文件是否有匹配
	if ip, exists := s.shieldManager.GetHostsIP(domain); exists {
		s.handleHostsResponse(w, r, ip)
		return
	}

	// 检查是否被屏蔽
	if s.shieldManager.IsBlocked(domain) {
		s.handleBlockedResponse(w, r, domain)
		return
	}

	// 转发到上游DNS服务器
	s.forwardDNSRequest(w, r, domain)
}

// handleHostsResponse 处理hosts文件匹配的响应
func (s *Server) handleHostsResponse(w dns.ResponseWriter, r *dns.Msg, ip string) {
	response := new(dns.Msg)
	response.SetReply(r)
	response.RecursionAvailable = true

	if len(r.Question) > 0 {
		q := r.Question[0]
		answer := new(dns.A)
		answer.Hdr = dns.RR_Header{
			Name:   q.Name,
			Rrtype: dns.TypeA,
			Class:  dns.ClassINET,
			Ttl:    300,
		}
		answer.A = net.ParseIP(ip)
		response.Answer = append(response.Answer, answer)
	}

	// 记录解析域名统计
	domain := ""
	if len(r.Question) > 0 {
		domain = r.Question[0].Name
		if len(domain) > 0 && domain[len(domain)-1] == '.' {
			domain = domain[:len(domain)-1]
		}
		s.updateResolvedDomainStats(domain)
	}

	w.WriteMsg(response)
	s.updateStats(func(stats *Stats) {
		stats.Allowed++
	})
}

// handleBlockedResponse 处理被屏蔽的域名响应
func (s *Server) handleBlockedResponse(w dns.ResponseWriter, r *dns.Msg, domain string) {
	logger.Info("域名被屏蔽", "domain", domain, "client", w.RemoteAddr())

	// 更新被屏蔽域名统计
	s.updateBlockedDomainStats(domain)

	// 更新总体统计
	s.updateStats(func(stats *Stats) {
		stats.Blocked++
	})

	response := new(dns.Msg)
	response.SetReply(r)
	response.RecursionAvailable = true

	// 获取屏蔽方法配置
	blockMethod := "NXDOMAIN" // 默认值
	customBlockIP := ""       // 默认值

	// 从Server结构体的shieldConfig字段获取配置
	if s.shieldConfig != nil {
		blockMethod = s.shieldConfig.BlockMethod
		customBlockIP = s.shieldConfig.CustomBlockIP
	}

	// 根据屏蔽方法返回不同的响应
	switch blockMethod {
	case "refused":
		// 返回拒绝查询响应
		response.SetRcode(r, dns.RcodeRefused)
	case "emptyIP":
		// 返回空IP响应
		if len(r.Question) > 0 && r.Question[0].Qtype == dns.TypeA {
			answer := new(dns.A)
			answer.Hdr = dns.RR_Header{
				Name:   r.Question[0].Name,
				Rrtype: dns.TypeA,
				Class:  dns.ClassINET,
				Ttl:    300,
			}
			answer.A = net.ParseIP("0.0.0.0") // 空IP
			response.Answer = append(response.Answer, answer)
		}
	case "customIP":
		// 返回自定义IP响应
		if len(r.Question) > 0 && r.Question[0].Qtype == dns.TypeA && customBlockIP != "" {
			answer := new(dns.A)
			answer.Hdr = dns.RR_Header{
				Name:   r.Question[0].Name,
				Rrtype: dns.TypeA,
				Class:  dns.ClassINET,
				Ttl:    300,
			}
			answer.A = net.ParseIP(customBlockIP)
			response.Answer = append(response.Answer, answer)
		}
	case "NXDOMAIN", "":
		fallthrough // 默认使用NXDOMAIN
	default:
		// 返回NXDOMAIN响应（域名不存在）
		response.SetRcode(r, dns.RcodeNameError)
	}

	w.WriteMsg(response)
	s.updateStats(func(stats *Stats) {
		stats.Blocked++
	})
}

// forwardDNSRequest 转发DNS请求到上游服务器
func (s *Server) forwardDNSRequest(w dns.ResponseWriter, r *dns.Msg, domain string) {
	// 尝试所有上游DNS服务器
	for _, upstream := range s.config.UpstreamDNS {
		response, rtt, err := s.resolver.Exchange(r, upstream)
		if err == nil && response != nil && response.Rcode == dns.RcodeSuccess {
			// 设置递归可用标志
			response.RecursionAvailable = true

			w.WriteMsg(response)
			logger.Debug("DNS查询成功", "domain", domain, "rtt", rtt, "server", upstream)

			// 记录解析域名统计
			s.updateResolvedDomainStats(domain)

			s.updateStats(func(stats *Stats) {
				stats.Allowed++
			})
			return
		}
	}

	// 所有上游服务器都失败，返回服务器失败错误
	response := new(dns.Msg)
	response.SetReply(r)
	response.RecursionAvailable = true
	response.SetRcode(r, dns.RcodeServerFailure)
	w.WriteMsg(response)

	logger.Error("DNS查询失败", "domain", domain)
	s.updateStats(func(stats *Stats) {
		stats.Errors++
	})
}

// updateBlockedDomainStats 更新被屏蔽域名统计
func (s *Server) updateBlockedDomainStats(domain string) {
	// 更新被屏蔽域名计数
	s.blockedDomainsMutex.Lock()
	defer s.blockedDomainsMutex.Unlock()

	if entry, exists := s.blockedDomains[domain]; exists {
		entry.Count++
		entry.LastSeen = time.Now()
	} else {
		s.blockedDomains[domain] = &BlockedDomain{
			Domain:   domain,
			Count:    1,
			LastSeen: time.Now(),
		}
	}

	// 更新统计数据
	now := time.Now()
	
	// 更新小时统计
	hourKey := now.Format("2006-01-02-15")
	s.hourlyStatsMutex.Lock()
	s.hourlyStats[hourKey]++
	s.hourlyStatsMutex.Unlock()
	
	// 更新每日统计
	dayKey := now.Format("2006-01-02")
	s.dailyStatsMutex.Lock()
	s.dailyStats[dayKey]++
	s.dailyStatsMutex.Unlock()
	
	// 更新每月统计
	monthKey := now.Format("2006-01")
	s.monthlyStatsMutex.Lock()
	s.monthlyStats[monthKey]++
	s.monthlyStatsMutex.Unlock()
}

// updateResolvedDomainStats 更新解析域名统计
func (s *Server) updateResolvedDomainStats(domain string) {
	s.resolvedDomainsMutex.Lock()
	defer s.resolvedDomainsMutex.Unlock()

	if entry, exists := s.resolvedDomains[domain]; exists {
		entry.Count++
		entry.LastSeen = time.Now()
	} else {
		s.resolvedDomains[domain] = &BlockedDomain{
			Domain:   domain,
			Count:    1,
			LastSeen: time.Now(),
		}
	}
}

// updateStats 更新统计信息
func (s *Server) updateStats(update func(*Stats)) {
	s.statsMutex.Lock()
	defer s.statsMutex.Unlock()
	update(s.stats)
}

// GetStats 获取DNS服务器统计信息
func (s *Server) GetStats() *Stats {
	s.statsMutex.Lock()
	defer s.statsMutex.Unlock()

	// 返回统计信息的副本
	return &Stats{
		Queries:   s.stats.Queries,
		Blocked:   s.stats.Blocked,
		Allowed:   s.stats.Allowed,
		Errors:    s.stats.Errors,
		LastQuery: s.stats.LastQuery,
	}
}

// GetTopBlockedDomains 获取TOP屏蔽域名列表
func (s *Server) GetTopBlockedDomains(limit int) []BlockedDomain {
	s.blockedDomainsMutex.RLock()
	defer s.blockedDomainsMutex.RUnlock()

	// 转换为切片
	domains := make([]BlockedDomain, 0, len(s.blockedDomains))
	for _, entry := range s.blockedDomains {
		domains = append(domains, *entry)
	}

	// 按计数排序
	sort.Slice(domains, func(i, j int) bool {
		return domains[i].Count > domains[j].Count
	})

	// 返回限制数量
	if len(domains) > limit {
		return domains[:limit]
	}
	return domains
}

// GetTopResolvedDomains 获取TOP解析域名
func (s *Server) GetTopResolvedDomains(limit int) []BlockedDomain {
	s.resolvedDomainsMutex.RLock()
	defer s.resolvedDomainsMutex.RUnlock()

	// 转换为切片
	domains := make([]BlockedDomain, 0, len(s.resolvedDomains))
	for _, entry := range s.resolvedDomains {
		domains = append(domains, *entry)
	}

	// 按数量排序
	sort.Slice(domains, func(i, j int) bool {
		return domains[i].Count > domains[j].Count
	})

	// 返回限制数量
	if len(domains) > limit {
		return domains[:limit]
	}
	return domains
}

// GetRecentBlockedDomains 获取最近屏蔽的域名列表
func (s *Server) GetRecentBlockedDomains(limit int) []BlockedDomain {
	s.blockedDomainsMutex.RLock()
	defer s.blockedDomainsMutex.RUnlock()

	// 转换为切片
	domains := make([]BlockedDomain, 0, len(s.blockedDomains))
	for _, entry := range s.blockedDomains {
		domains = append(domains, *entry)
	}

	// 按时间排序
	sort.Slice(domains, func(i, j int) bool {
		return domains[i].LastSeen.After(domains[j].LastSeen)
	})

	// 返回限制数量
	if len(domains) > limit {
		return domains[:limit]
	}
	return domains
}

// GetHourlyStats 获取每小时统计数据
func (s *Server) GetHourlyStats() map[string]int64 {
	s.hourlyStatsMutex.RLock()
	defer s.hourlyStatsMutex.RUnlock()

	// 返回副本
	result := make(map[string]int64)
	for k, v := range s.hourlyStats {
		result[k] = v
	}
	return result
}

// GetDailyStats 获取每日统计数据
func (s *Server) GetDailyStats() map[string]int64 {
	s.dailyStatsMutex.RLock()
	defer s.dailyStatsMutex.RUnlock()

	// 返回副本
	result := make(map[string]int64)
	for k, v := range s.dailyStats {
		result[k] = v
	}
	return result
}

// GetMonthlyStats 获取每月统计数据
func (s *Server) GetMonthlyStats() map[string]int64 {
	s.monthlyStatsMutex.RLock()
	defer s.monthlyStatsMutex.RUnlock()

	// 返回副本
	result := make(map[string]int64)
	for k, v := range s.monthlyStats {
		result[k] = v
	}
	return result
}

// loadStatsData 从文件加载统计数据
func (s *Server) loadStatsData() {
	if s.config.StatsFile == "" {
		return
	}

	// 检查文件是否存在
	data, err := ioutil.ReadFile(s.config.StatsFile)
	if err != nil {
		if !os.IsNotExist(err) {
			logger.Error("读取统计数据文件失败", "error", err)
		}
		return
	}

	var statsData StatsData
	err = json.Unmarshal(data, &statsData)
	if err != nil {
		logger.Error("解析统计数据失败", "error", err)
		return
	}

	// 恢复统计数据
	s.statsMutex.Lock()
	if statsData.Stats != nil {
		s.stats = statsData.Stats
	}
	s.statsMutex.Unlock()

	s.blockedDomainsMutex.Lock()
	if statsData.BlockedDomains != nil {
		s.blockedDomains = statsData.BlockedDomains
	}
	s.blockedDomainsMutex.Unlock()

	s.resolvedDomainsMutex.Lock()
	if statsData.ResolvedDomains != nil {
		s.resolvedDomains = statsData.ResolvedDomains
	}
	s.resolvedDomainsMutex.Unlock()

	s.hourlyStatsMutex.Lock()
	if statsData.HourlyStats != nil {
		s.hourlyStats = statsData.HourlyStats
	}
	s.hourlyStatsMutex.Unlock()
	
	s.dailyStatsMutex.Lock()
	if statsData.DailyStats != nil {
		s.dailyStats = statsData.DailyStats
	}
	s.dailyStatsMutex.Unlock()
	
	s.monthlyStatsMutex.Lock()
	if statsData.MonthlyStats != nil {
		s.monthlyStats = statsData.MonthlyStats
	}
	s.monthlyStatsMutex.Unlock()

	logger.Info("统计数据加载成功")
}

// saveStatsData 保存统计数据到文件
func (s *Server) saveStatsData() {
	if s.config.StatsFile == "" {
		return
	}

	// 创建数据目录
	statsDir := filepath.Dir(s.config.StatsFile)
	err := os.MkdirAll(statsDir, 0755)
	if err != nil {
		logger.Error("创建统计数据目录失败", "error", err)
		return
	}

	// 收集所有统计数据
	statsData := &StatsData{
		Stats:       s.GetStats(),
		LastSaved:   time.Now(),
	}

	// 复制域名数据
	s.blockedDomainsMutex.RLock()
	statsData.BlockedDomains = make(map[string]*BlockedDomain)
	for k, v := range s.blockedDomains {
		statsData.BlockedDomains[k] = v
	}
	s.blockedDomainsMutex.RUnlock()

	s.resolvedDomainsMutex.RLock()
	statsData.ResolvedDomains = make(map[string]*BlockedDomain)
	for k, v := range s.resolvedDomains {
		statsData.ResolvedDomains[k] = v
	}
	s.resolvedDomainsMutex.RUnlock()

	s.hourlyStatsMutex.RLock()
	statsData.HourlyStats = make(map[string]int64)
	for k, v := range s.hourlyStats {
		statsData.HourlyStats[k] = v
	}
	s.hourlyStatsMutex.RUnlock()
	
	s.dailyStatsMutex.RLock()
	statsData.DailyStats = make(map[string]int64)
	for k, v := range s.dailyStats {
		statsData.DailyStats[k] = v
	}
	s.dailyStatsMutex.RUnlock()
	
	s.monthlyStatsMutex.RLock()
	statsData.MonthlyStats = make(map[string]int64)
	for k, v := range s.monthlyStats {
		statsData.MonthlyStats[k] = v
	}
	s.monthlyStatsMutex.RUnlock()

	// 序列化数据
	jsonData, err := json.MarshalIndent(statsData, "", "  ")
	if err != nil {
		logger.Error("序列化统计数据失败", "error", err)
		return
	}

	// 写入文件
	err = ioutil.WriteFile(s.config.StatsFile, jsonData, 0644)
	if err != nil {
		logger.Error("保存统计数据到文件失败", "error", err)
		return
	}

	logger.Info("统计数据保存成功")
}

// startAutoSave 启动自动保存功能
func (s *Server) startAutoSave() {
	if s.config.StatsFile == "" || s.config.SaveInterval <= 0 {
		return
	}

	// 设置定时器
	s.saveTicker = time.NewTicker(time.Duration(s.config.SaveInterval) * time.Second)
	defer s.saveTicker.Stop()

	logger.Info("启动统计数据自动保存功能", "interval", s.config.SaveInterval, "file", s.config.StatsFile)

	// 定期保存数据
	for {
		select {
		case <-s.saveTicker.C:
			s.saveStatsData()
		case <-s.saveDone:
			return
		}
	}
}
