# -*- coding: UTF-8 -*-
'''
@Project :domainScanDemo
@File    :domain_collector.py
@IDE     :PyCharm
@Author  :梦伴
@Date    :2026/4/8 23:50
@explain : 域名收集器
'''

import re
import tldextract
from loguru import logger
from app.utils.database import Database
from app.utils.domain_utils import normalize_domain
from app.config import config


class DomainCollector:
    """
    域名收集器
    """
    
    def __init__(self):
        """
        初始化域名收集器
        """
        self.db = Database(
            host=config.DB_HOST,
            port=config.DB_PORT,
            database=config.DB_DATABASE,
            user=config.DB_USER,
            password=config.DB_PASSWORD
        )
    
    def add_domain(self, domain, source_type):
        """
        添加域名
        
        :param domain: 域名
        :param source_type: 来源类型
        :return: bool - 是否添加成功
        """
        try:
            # 标准化域名
            normalized_domain = normalize_domain(domain)
            if not normalized_domain:
                logger.warning(f"无效域名: {domain}")
                return False
            
            # 提取顶级域名
            ext = tldextract.extract(normalized_domain)
            tld = ext.suffix
            
            # 只保留 .com 和 .net
            if tld not in ['com', 'net']:
                logger.warning(f"不支持的顶级域名: {tld}")
                return False
            
            # 检查是否已存在
            if self.db.domain_exists(normalized_domain):
                logger.info(f"域名已存在: {normalized_domain}")
                return False
            
            # 添加域名
            domain_id = self.db.add_domain(normalized_domain, tld, source_type)
            if domain_id:
                # 创建检测任务
                self.db.create_detect_task(domain_id, 1)  # 1 表示基础检测
                logger.info(f"成功添加域名: {normalized_domain}")
                return True
            else:
                logger.error(f"添加域名失败: {normalized_domain}")
                return False
        except Exception as e:
            logger.error(f"添加域名出错: {e}")
            return False
    
    def add_domains_batch(self, domains, source_type, batch_size=1000, dry_run=False):
        """
        批量添加域名
        
        :param domains: 域名列表
        :param source_type: 来源类型
        :param batch_size: 批量大小
        :param dry_run: 是否仅进行干运行（不实际添加域名）
        :return: dict - 统计信息
        """
        try:
            # 统计信息
            stats = {
                'total': len(domains),
                'valid': 0,
                'added': 0,
                'exists': 0,
                'invalid': 0,
                'failed': 0
            }
            
            # 处理大规模数据时，分批进行标准化和过滤
            normalized_domains = []
            batch_domains = []
            
            for i, domain in enumerate(domains):
                normalized = normalize_domain(domain)
                if not normalized:
                    stats['invalid'] += 1
                    continue
                
                # 提取顶级域名
                ext = tldextract.extract(normalized)
                tld = ext.suffix
                
                # 只保留 .com 和 .net
                if tld not in ['com', 'net']:
                    stats['invalid'] += 1
                    continue
                
                normalized_domains.append((normalized, tld))
                batch_domains.append(normalized)
                stats['valid'] += 1
                
                # 每1000个域名检查一次，避免内存占用过高
                if (i + 1) % 1000 == 0:
                    logger.info(f"已处理 {i + 1}/{len(domains)} 个域名")
            
            logger.info(f"域名标准化完成，有效域名: {stats['valid']}")
            
            # 提取所有域名
            all_domains = [domain for domain, tld in normalized_domains]
            
            # 批量检查域名是否存在
            existing_domains = self.db.check_domains_exist(all_domains)
            existing_set = set(existing_domains)
            
            # 准备批量添加数据
            batch_data = []
            for domain, tld in normalized_domains:
                if domain not in existing_set:
                    batch_data.append((domain, tld, source_type))
            
            stats['exists'] = len(existing_domains)
            stats['valid'] = len(normalized_domains)
            
            # 干运行模式下直接返回统计信息
            if dry_run:
                stats['added'] = len(batch_data)
                logger.info(f"干运行模式：准备添加 {len(batch_data)} 个新域名")
                return stats
            
            logger.info(f"准备添加 {len(batch_data)} 个新域名")
            
            # 分批次添加
            for i in range(0, len(batch_data), batch_size):
                batch = batch_data[i:i+batch_size]
                added_count = self.db.add_domains_batch(batch)
                stats['added'] += added_count
                stats['failed'] += len(batch) - added_count
                
                # 每处理一批，记录一次进度
                if (i + len(batch)) % (batch_size * 10) == 0:
                    logger.info(f"已添加 {i + len(batch)}/{len(batch_data)} 个域名")
            
            logger.info(f"批量添加域名完成: {stats}")
            return stats
        except Exception as e:
            logger.error(f"批量添加域名出错: {e}")
            return {
                'total': len(domains),
                'valid': 0,
                'added': 0,
                'exists': 0,
                'invalid': 0,
                'failed': len(domains)
            }
    
    def import_from_file(self, file_path, source_type, batch_size=1000):
        """
        从文件导入域名
        
        :param file_path: 文件路径
        :param source_type: 来源类型
        :param batch_size: 批量大小
        :return: dict - 导入统计信息
        """
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                domains = f.readlines()
            
            # 提取域名
            domain_list = []
            for domain in domains:
                domain = domain.strip()
                if domain:
                    domain_list.append(domain)
            
            # 批量添加域名
            stats = self.add_domains_batch(domain_list, source_type, batch_size)
            
            logger.info(f"从文件导入完成: {stats}")
            return stats
        except Exception as e:
            logger.error(f"从文件导入出错: {e}")
            return {
                'total': 0,
                'valid': 0,
                'added': 0,
                'exists': 0,
                'invalid': 0,
                'failed': 0
            }
    
    def collect_from_juming(self, type_='一口价'):
        """
        从聚名网收集域名
        
        :param type_: 类型，一口价或过期删除
        :return: int - 收集到的域名数量
        """
        # 这里可以添加从聚名网收集域名的逻辑
        logger.info(f"从聚名网收集 {type_} 域名")
        # 模拟收集结果
        return 0
    
    def collect_from_search_engine(self, keyword, limit=100):
        """
        从搜索引擎收集域名
        
        :param keyword: 关键词
        :param limit: 限制数量
        :return: int - 收集到的域名数量
        """
        # 这里可以添加从搜索引擎收集域名的逻辑
        logger.info(f"从搜索引擎收集域名，关键词: {keyword}, 限制: {limit}")
        # 模拟收集结果
        return 0
    
    def collect_from_enterprise_directory(self, url, limit=100):
        """
        从企业目录收集域名
        
        :param url: 企业目录URL
        :param limit: 限制数量
        :return: int - 收集到的域名数量
        """
        # 这里可以添加从企业目录收集域名的逻辑
        logger.info(f"从企业目录收集域名，URL: {url}, 限制: {limit}")
        # 模拟收集结果
        return 0
    
    def collect_from_zone_file(self, file_path):
        """
        从Zone File收集域名
        
        :param file_path: Zone File路径
        :return: int - 收集到的域名数量
        """
        # 这里可以添加从Zone File收集域名的逻辑
        logger.info(f"从Zone File收集域名，文件: {file_path}")
        # 模拟收集结果
        return 0