import requests
import pandas as pd
import time
import os
import random

# ============================================================
# 配置
# ============================================================

START_DATE = "2020-01-01"

# 你真正需要的数据截止日期
END_DATE = "2026-08-31"

# 实际请求时，向后多请求几天
# 例如 END_DATE=2026-08-31
# 实际请求到 2026-09-03
REQUEST_EXTRA_DAYS = 3

# 输出目录
OUTPUT_DIR = "weather_data"
os.makedirs(OUTPUT_DIR, exist_ok=True)

# 最终合并文件
FINAL_FILE = "china_city_temperature_2020_2026.csv"

# Open-Meteo 历史 API
URL = "https://archive-api.open-meteo.com/v1/archive"


# ============================================================
# 城市
# ============================================================

cities = {
    "北京": (39.9042, 116.4074),
    "上海": (31.2304, 121.4737),
    "广州": (23.1291, 113.2644),
    "深圳": (22.5431, 114.0579),
    "成都": (30.5728, 104.0668),
    "武汉": (30.5928, 114.3055),
    "西安": (34.3416, 108.9398),
    "南京": (32.0603, 118.7969),
    "青岛": (36.0671, 120.3826),
    "哈尔滨": (45.8038, 126.5349),
    "乌鲁木齐": (43.8256, 87.6168),
    "昆明": (25.0389, 102.7183),
}


# ============================================================
# 计算实际请求结束日期
# ============================================================

request_end_date = (
    pd.Timestamp(END_DATE)
    + pd.Timedelta(days=REQUEST_EXTRA_DAYS)
).strftime("%Y-%m-%d")

print("=" * 60)
print("数据配置")
print("=" * 60)
print(f"实际需要：{START_DATE} ~ {END_DATE}")
print(f"API请求： {START_DATE} ~ {request_end_date}")
print(f"额外请求：{REQUEST_EXTRA_DAYS} 天")
print("=" * 60)


# ============================================================
# 创建 Session
# ============================================================

session = requests.Session()

session.headers.update({
    "User-Agent": "Mozilla/5.0 (Weather Data Collector)"
})


# ============================================================
# 获取一个城市
# ============================================================

def get_city_weather(city, lat, lon, max_retries=6):

    output_file = os.path.join(
        OUTPUT_DIR,
        f"{city}_{START_DATE}_{END_DATE}.csv"
    )

    # --------------------------------------------------------
    # 计算真正需要的数据天数
    # --------------------------------------------------------

    expected_days = (
        pd.Timestamp(END_DATE)
        - pd.Timestamp(START_DATE)
    ).days + 1

    # --------------------------------------------------------
    # 如果已经下载成功，直接跳过
    #
    # 注意：
    # 不再简单使用 len(old) > 3000 判断。
    # 必须同时检查：
    # 1. 日期数量
    # 2. 三个温度字段没有缺失
    # --------------------------------------------------------

    if os.path.exists(output_file):

        try:

            old = pd.read_csv(output_file)

            required_columns = {
                "date",
                "city",
                "mean_temp",
                "max_temp",
                "min_temp"
            }

            columns_ok = required_columns.issubset(
                old.columns
            )

            row_count_ok = len(old) == expected_days

            temp_ok = False

            if columns_ok:

                temp_ok = not old[
                    [
                        "mean_temp",
                        "max_temp",
                        "min_temp"
                    ]
                ].isna().any().any()

            if (
                columns_ok
                and row_count_ok
                and temp_ok
            ):

                print(
                    f"✓ {city} 已存在且完整，跳过"
                )

                return old

            else:

                print(
                    f"⚠ {city} 已存在，但数据不完整，重新获取"
                )

        except Exception as e:

            print(
                f"⚠ {city} 原文件读取失败，重新获取：{e}"
            )


    # ========================================================
    # Open-Meteo 请求参数
    # ========================================================

    params = {

        "latitude": lat,
        "longitude": lon,

        # ----------------------------------------------------
        # 重点：
        # START_DATE 不变
        # END_DATE 向后增加 3 天
        # ----------------------------------------------------

        "start_date": START_DATE,
        "end_date": request_end_date,

        "daily": (
            "temperature_2m_mean,"
            "temperature_2m_max,"
            "temperature_2m_min"
        ),

        "timezone": "Asia/Shanghai",

        # 长期历史数据
        "models": "era5_land",

        # 陆地网格选择
        "cell_selection": "land",

        # 摄氏度
        "temperature_unit": "celsius",
    }


    # ========================================================
    # 开始请求
    # ========================================================

    for attempt in range(1, max_retries + 1):

        try:

            print(
                f"正在获取：{city} "
                f"(第 {attempt}/{max_retries} 次)"
            )

            print(
                f"  API范围：{START_DATE} ~ {request_end_date}"
            )

            response = session.get(
                URL,
                params=params,
                timeout=(15, 180)
            )

            response.raise_for_status()

            data = response.json()

            if data.get("error"):

                raise RuntimeError(
                    data.get(
                        "reason",
                        "Open-Meteo API error"
                    )
                )

            daily = data["daily"]


            # =================================================
            # 构造 DataFrame
            # =================================================

            df = pd.DataFrame({

                "date": daily["time"],

                "city": city,

                "mean_temp": daily[
                    "temperature_2m_mean"
                ],

                "max_temp": daily[
                    "temperature_2m_max"
                ],

                "min_temp": daily[
                    "temperature_2m_min"
                ],
            })


            # =================================================
            # 重点：
            # 只保留用户真正要求的 END_DATE
            # =================================================

            df["date"] = pd.to_datetime(df["date"])

            target_end = pd.Timestamp(END_DATE)

            df = df[
                df["date"] <= target_end
            ].copy()


            # 日期重新排序
            df = df.sort_values(
                "date"
            ).reset_index(drop=True)


            # =================================================
            # 数据数量检查
            # =================================================

            if len(df) != expected_days:

                raise RuntimeError(
                    f"数据数量异常：{len(df)} 条，"
                    f"预期 {expected_days} 条"
                )


            # =================================================
            # 日期连续性检查
            # =================================================

            expected_dates = pd.date_range(
                start=START_DATE,
                end=END_DATE,
                freq="D"
            )

            actual_dates = pd.DatetimeIndex(
                df["date"]
            )

            if not actual_dates.equals(
                expected_dates
            ):

                raise RuntimeError(
                    "日期不连续或存在缺失日期"
                )


            # =================================================
            # 温度数据缺失检查
            # =================================================

            temperature_columns = [
                "mean_temp",
                "max_temp",
                "min_temp"
            ]

            missing_count = (
                df[temperature_columns]
                .isna()
                .any(axis=1)
                .sum()
            )


            if missing_count > 0:

                missing_dates = df.loc[
                    df[temperature_columns]
                    .isna()
                    .any(axis=1),
                    "date"
                ].dt.strftime(
                    "%Y-%m-%d"
                ).tolist()

                raise RuntimeError(
                    f"存在 {missing_count} 天温度数据为空："
                    f"{missing_dates[-10:]}"
                )


            # =================================================
            # 保存单城市文件
            # =================================================

            df["date"] = df[
                "date"
            ].dt.strftime("%Y-%m-%d")

            df.to_csv(
                output_file,
                index=False,
                encoding="utf-8-sig"
            )


            print(
                f"✓ {city} 获取成功："
                f"{len(df)} 天"
            )

            print(
                f"  数据范围："
                f"{df['date'].iloc[0]} ~ "
                f"{df['date'].iloc[-1]}"
            )

            return df


        # =====================================================
        # 请求失败
        # =====================================================

        except Exception as e:

            print(
                f"✗ {city} 第 {attempt} 次失败："
                f"{type(e).__name__}: {e}"
            )

            if attempt < max_retries:

                # 指数退避
                wait = min(
                    60,
                    5 * (2 ** (attempt - 1))
                )

                # 随机延迟
                wait += random.uniform(0, 3)

                print(
                    f"等待 {wait:.1f} 秒后重试..."
                )

                time.sleep(wait)

            else:

                print(
                    f"✗ {city} 最终失败，"
                    f"后面会继续处理其他城市"
                )

                return None


# ============================================================
# 开始下载
# ============================================================

all_data = []

failed = []


for city, (lat, lon) in cities.items():

    df = get_city_weather(
        city,
        lat,
        lon
    )

    if df is not None:

        all_data.append(df)

    else:

        failed.append(city)

    # 城市之间稍微等待
    time.sleep(2)


# ============================================================
# 合并所有成功数据
# ============================================================

if all_data:

    result = pd.concat(
        all_data,
        ignore_index=True
    )

    result.to_csv(
        FINAL_FILE,
        index=False,
        encoding="utf-8-sig"
    )

    print("\n" + "=" * 60)
    print("下载完成")
    print("=" * 60)

    print(f"成功城市：{len(all_data)}")
    print(f"失败城市：{len(failed)}")

    print(f"总数据量：{len(result):,}")

    print(
        f"理论数据量："
        f"{expected_days * len(cities):,}"
    )

    print(
        f"最终文件：{FINAL_FILE}"
    )

    if failed:

        print("\n失败城市：")

        for city in failed:
            print("  -", city)


else:

    print(
        "没有成功获取任何城市数据"
    )


# ============================================================
# 数据检查
# ============================================================

print("\n数据预览：")

if all_data:

    print(
        result.head(10)
    )

    print("\n每个城市数据量：")

    print(
        result.groupby("city").size()
    )

    print("\n缺失值检查：")

    print(
        result.isna().sum()
    )

    print("\n日期范围检查：")

    print(
        result.groupby("city")["date"]
        .agg(["min", "max", "count"])
    )