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

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

START_DATE = "2020-01-01"
END_DATE   = "2026-08-31"

# 输出目录
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),
}

# ============================================================
# 创建 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"
    )

    # --------------------------------------------------------
    # 如果已经下载成功，直接跳过
    # --------------------------------------------------------

    if os.path.exists(output_file):

        try:
            old = pd.read_csv(output_file)

            if len(old) > 3000:
                print(f"✓ {city} 已存在，跳过")
                return old

        except Exception:
            pass

    params = {
        "latitude": lat,
        "longitude": lon,

        "start_date": START_DATE,
        "end_date": END_DATE,

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

        "timezone": "Asia/Shanghai",

        # 长期历史数据使用 ERA5-Land
        "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} 次)"
            )

            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"]

            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"],
            })

            # ------------------------------------------------
            # 检查数据完整性
            # ------------------------------------------------

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

            if len(df) < expected_days * 0.99:

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

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

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

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

            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"最终文件：{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())