main.py

# -*- coding: utf-8 -*-
"""
天气爬虫主程序
支持命令行参数和配置文件两种方式运行
"""

import argparse
import asyncio
import logging
import os
import sys
from datetime import datetime, timedelta

# 添加当前目录到路径
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from config import (
    CITY_CODES,
    DATA_FIELDS,
    OUTPUT_FORMAT,
    OUTPUT_DIR,
    LOG_LEVEL,
    LOG_TO_FILE,
    LOG_FILE,
    DEFAULT_DAYS,
)
from crawler import WeatherCrawler, CustomWeatherCrawler
from exporter import DataExporter


def setup_logging(level: str = LOG_LEVEL, to_file: bool = LOG_TO_FILE):
    """
    配置日志

    Args:
        level: 日志级别
        to_file: 是否输出到文件
    """
    log_format = "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
    handlers = [logging.StreamHandler()]

    if to_file:
        log_dir = os.path.dirname(LOG_FILE)
        if log_dir:
            os.makedirs(log_dir, exist_ok=True)
        handlers.append(logging.FileHandler(LOG_FILE, encoding='utf-8'))

    logging.basicConfig(
        level=getattr(logging, level.upper()),
        format=log_format,
        handlers=handlers,
    )


def parse_args():
    """
    解析命令行参数

    Returns:
        解析后的参数
    """
    parser = argparse.ArgumentParser(
        description="天气数据爬虫 - 基于crawl4ai",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例用法:
  # 爬取北京最近7天的天气
  python main.py -c 北京

  # 爬取多个城市指定日期范围的天气
  python main.py -c 北京 上海 -s 2024-01-01 -e 2024-01-31

  # 导出为Excel格式
  python main.py -c 广州 -f xlsx

  # 指定输出目录
  python main.py -c 深圳 -o ./data

可用城市: """ + ", ".join(CITY_CODES.keys())
    )

    parser.add_argument(
        "-c", "--cities",
        nargs="+",
        required=True,
        help="要爬取的城市列表 (必填)"
    )

    parser.add_argument(
        "-s", "--start-date",
        type=str,
        default=None,
        help="开始日期 (格式: YYYY-MM-DD)"
    )

    parser.add_argument(
        "-e", "--end-date",
        type=str,
        default=None,
        help="结束日期 (格式: YYYY-MM-DD)"
    )

    parser.add_argument(
        "-f", "--format",
        type=str,
        choices=["csv", "xlsx"],
        default=OUTPUT_FORMAT,
        help=f"输出格式 (默认: {OUTPUT_FORMAT})"
    )

    parser.add_argument(
        "-o", "--output",
        type=str,
        default=OUTPUT_DIR,
        help=f"输出目录 (默认: {OUTPUT_DIR})"
    )

    parser.add_argument(
        "--fields",
        nargs="+",
        default=None,
        help="要爬取的字段列表 (可选: date, weather, temp_high, temp_low, wind_direction, wind_level)"
    )

    parser.add_argument(
        "-v", "--verbose",
        action="store_true",
        help="显示详细日志"
    )

    parser.add_argument(
        "-m", "--monthly",
        action="store_true",
        help="爬取月度统计数据（平均温度、极端温度、空气质量等）"
    )

    parser.add_argument(
        "--daily",
        action="store_true",
        default=True,
        help="爬取每日天气数据（默认开启）"
    )

    return parser.parse_args()


def validate_date(date_str: str) -> bool:
    """
    验证日期格式

    Args:
        date_str: 日期字符串

    Returns:
        是否有效
    """
    try:
        datetime.strptime(date_str, "%Y-%m-%d")
        return True
    except ValueError:
        return False


def get_default_date_range():
    """
    获取默认日期范围 (最近N天)

    Returns:
        (start_date, end_date) 元组
    """
    end_date = datetime.now()
    start_date = end_date - timedelta(days=DEFAULT_DAYS)
    return start_date.strftime("%Y-%m-%d"), end_date.strftime("%Y-%m-%d")


async def run_crawler(args):
    """
    运行爬虫

    Args:
        args: 命令行参数
    """
    logger = logging.getLogger(__name__)

    # 验证城市
    invalid_cities = [c for c in args.cities if c not in CITY_CODES]
    if invalid_cities:
        logger.warning(f"以下城市不在预设列表中，将尝试直接使用: {invalid_cities}")

    # 处理日期
    start_date = args.start_date
    end_date = args.end_date

    if not start_date or not end_date:
        start_date, end_date = get_default_date_range()
        logger.info(f"使用默认日期范围: {start_date} 至 {end_date}")
    else:
        if not validate_date(start_date) or not validate_date(end_date):
            logger.error("日期格式错误，请使用 YYYY-MM-DD 格式")
            return

    # 处理字段配置
    fields = DATA_FIELDS.copy()
    if args.fields:
        # 只启用指定的字段
        fields = {k: (k in args.fields) for k in DATA_FIELDS.keys()}
        fields["city"] = True  # 城市字段始终保留

    # 创建爬虫
    crawler = WeatherCrawler(
        cities=args.cities,
        start_date=start_date,
        end_date=end_date,
        fields=fields,
    )

    # 执行爬取
    logger.info("=" * 50)
    logger.info("开始爬取天气数据")
    logger.info(f"城市: {args.cities}")
    logger.info(f"日期范围: {start_date} 至 {end_date}")
    logger.info(f"输出格式: {args.format}")
    logger.info(f"爬取模式: {'月度统计' if args.monthly else '每日数据'}")
    logger.info("=" * 50)

    try:
        exporter = DataExporter(
            output_format=args.format,
            output_dir=args.output,
        )
        city_name = "_".join(args.cities) if len(args.cities) <= 3 else "multiple"

        # 爬取月度统计数据
        if args.monthly:
            monthly_results = await crawler.crawl_monthly_stats()
            if monthly_results:
                filepath = exporter.export(monthly_results, city=f"{city_name}_monthly")
                if filepath:
                    logger.info(f"月度统计数据已导出到: {filepath}")
                    logger.info(f"共导出 {len(monthly_results)} 条月度记录")
            else:
                logger.warning("未获取到月度统计数据")

        # 爬取每日数据（默认行为，除非只指定了-m）
        if not args.monthly or args.daily:
            if not args.monthly:  # 只有不是月度模式时才爬取每日数据
                results = await crawler.crawl()
                if results:
                    filepath = exporter.export(results, city=city_name)
                    if filepath:
                        logger.info(f"每日数据已导出到: {filepath}")
                        logger.info(f"共导出 {len(results)} 条记录")
                else:
                    logger.warning("未获取到每日数据")

    except Exception as e:
        logger.error(f"爬取过程中发生错误: {e}")
        raise


def main():
    """主函数入口"""
    args = parse_args()

    # 配置日志
    log_level = "DEBUG" if args.verbose else LOG_LEVEL
    setup_logging(level=log_level)

    # 运行爬虫
    asyncio.run(run_crawler(args))


if __name__ == "__main__":
    main()