#!/usr/bin/env python3
"""Count NOAA CFS reforecast GRIB2 files and create a Word report.

Only THREDDS catalog metadata are downloaded; the GRIB2 payloads themselves are
not downloaded.  Daily results are cached in the working directory so the run
can be resumed safely.
"""

from __future__ import annotations

import argparse
import csv
import json
import re
import threading
import time
import urllib.error
import urllib.request
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import date, timedelta
from pathlib import Path


ROOT = Path(__file__).resolve().parent
CACHE_FILE = ROOT / "noaa_grib2_daily_counts.json"
CSV_FILE = ROOT / "noaa_grib2_monthly_counts.csv"
DOCX_FILE = ROOT / "thong_ke_so_luong_file_grib2_flx_pgb.docx"
YEARS = range(1999, 2012)
DATASETS = {
    "FLX": "model-cfs_refor_6h_90d_flx",
    "PGB": "model-cfs_refor_6h_90d_pgb",
}
CATALOG_ROOT = "https://www.ncei.noaa.gov/thredds/catalog"
GRIB2_URLPATH = re.compile(rb'urlPath="[^"]+\.grb2"', re.IGNORECASE)
PRINT_LOCK = threading.Lock()


def all_days() -> list[date]:
    result: list[date] = []
    current = date(1999, 1, 1)
    end = date(2011, 3, 31)
    while current <= end:
        result.append(current)
        current += timedelta(days=1)
    return result


def catalog_url(dataset_path: str, day: date) -> str:
    stamp = day.strftime("%Y%m%d")
    return (
        f"{CATALOG_ROOT}/{dataset_path}/{day:%Y}/{day:%Y%m}/"
        f"{stamp}/catalog.xml"
    )


def fetch_count(label: str, dataset_path: str, day: date) -> tuple[str, str, int]:
    url = catalog_url(dataset_path, day)
    last_error: Exception | None = None
    for attempt in range(1, 7):
        try:
            request = urllib.request.Request(
                url,
                headers={
                    "User-Agent": "NOAA-GRIB2-statistics/1.0",
                    "Accept": "application/xml",
                },
            )
            with urllib.request.urlopen(request, timeout=60) as response:
                body = response.read()
            if b"<catalog" not in body:
                raise ValueError("response is not a THREDDS catalog")
            paths = GRIB2_URLPATH.findall(body)
            return label, day.isoformat(), len(set(paths))
        except (urllib.error.URLError, TimeoutError, ValueError) as exc:
            if isinstance(exc, urllib.error.HTTPError) and exc.code == 404:
                return label, day.isoformat(), 0
            last_error = exc
            time.sleep(min(2 ** (attempt - 1), 12))
    raise RuntimeError(f"Failed after retries: {url}: {last_error}")


def load_cache() -> dict[str, dict[str, int]]:
    if not CACHE_FILE.exists():
        return {label: {} for label in DATASETS}
    with CACHE_FILE.open("r", encoding="utf-8") as handle:
        data = json.load(handle)
    return {label: dict(data.get(label, {})) for label in DATASETS}


def save_cache(cache: dict[str, dict[str, int]]) -> None:
    temporary = CACHE_FILE.with_suffix(".json.tmp")
    with temporary.open("w", encoding="utf-8") as handle:
        json.dump(cache, handle, ensure_ascii=False, indent=2, sort_keys=True)
    temporary.replace(CACHE_FILE)


def collect(workers: int) -> dict[str, dict[str, int]]:
    cache = load_cache()
    pending = [
        (label, path, day)
        for label, path in DATASETS.items()
        for day in all_days()
        if day.isoformat() not in cache[label]
    ]
    total = len(pending)
    if not pending:
        print("All daily catalogs are already cached.")
        return cache

    print(f"Reading {total:,} daily catalogs with {workers} workers...")
    errors: list[str] = []
    completed = 0
    with ThreadPoolExecutor(max_workers=workers) as executor:
        futures = {
            executor.submit(fetch_count, label, path, day): (label, day)
            for label, path, day in pending
        }
        for future in as_completed(futures):
            completed += 1
            label, day = futures[future]
            try:
                result_label, stamp, count = future.result()
                cache[result_label][stamp] = count
            except Exception as exc:  # Keep collecting other dates for resumption.
                errors.append(f"{label} {day.isoformat()}: {exc}")
            if completed % 100 == 0 or completed == total:
                save_cache(cache)
                with PRINT_LOCK:
                    print(f"  {completed:,}/{total:,} complete; errors={len(errors)}")

    save_cache(cache)
    if errors:
        error_file = ROOT / "noaa_grib2_errors.txt"
        error_file.write_text("\n".join(errors) + "\n", encoding="utf-8")
        raise RuntimeError(
            f"{len(errors)} catalogs failed; details in {error_file.name}. "
            "Run the script again to retry only missing dates."
        )
    return cache


def monthly_counts(cache: dict[str, dict[str, int]]) -> dict[str, dict[int, list[int]]]:
    expected = {day.isoformat() for day in all_days()}
    for label in DATASETS:
        missing = expected.difference(cache[label])
        if missing:
            raise ValueError(f"{label} is missing {len(missing)} daily catalogs")

    result: dict[str, dict[int, list[int]]] = {}
    for label in DATASETS:
        by_month: dict[tuple[int, int], int] = defaultdict(int)
        for stamp, count in cache[label].items():
            day = date.fromisoformat(stamp)
            by_month[(day.year, day.month)] += int(count)
        result[label] = {
            year: [by_month[(year, month)] for month in range(1, 13)]
            for year in YEARS
        }
    return result


def write_csv(monthly: dict[str, dict[int, list[int]]]) -> None:
    with CSV_FILE.open("w", newline="", encoding="utf-8-sig") as handle:
        writer = csv.writer(handle)
        writer.writerow(["Bộ dữ liệu", "Năm", *[f"Tháng {m}" for m in range(1, 13)]])
        for label in DATASETS:
            for year in YEARS:
                writer.writerow([label, year, *monthly[label][year]])


def set_cell_text(cell, value: object, *, bold: bool = False, size: int = 8) -> None:
    from docx.shared import Pt

    cell.text = ""
    paragraph = cell.paragraphs[0]
    paragraph.alignment = 1
    run = paragraph.add_run(str(value))
    run.bold = bold
    run.font.name = "Arial"
    run.font.size = Pt(size)


def shade_cell(cell, fill: str) -> None:
    from docx.oxml import OxmlElement
    from docx.oxml.ns import qn

    shading = OxmlElement("w:shd")
    shading.set(qn("w:fill"), fill)
    cell._tc.get_or_add_tcPr().append(shading)


def add_table(document, title: str, rows: dict[int, list[int]]) -> None:
    from docx.enum.text import WD_ALIGN_PARAGRAPH

    heading = document.add_paragraph()
    heading.alignment = WD_ALIGN_PARAGRAPH.LEFT
    run = heading.add_run(title)
    run.bold = True
    run.font.name = "Arial"

    table = document.add_table(rows=1, cols=13)
    table.style = "Table Grid"
    table.autofit = False
    headers = ["Năm", *[f"Tháng {month}" for month in range(1, 13)]]
    for index, header in enumerate(headers):
        set_cell_text(table.rows[0].cells[index], header, bold=True, size=7)
        shade_cell(table.rows[0].cells[index], "D9EAF7")
    for year in YEARS:
        cells = table.add_row().cells
        set_cell_text(cells[0], year, bold=True)
        for month_index, count in enumerate(rows[year], start=1):
            set_cell_text(cells[month_index], count)
    document.add_paragraph()


def write_docx(monthly: dict[str, dict[int, list[int]]]) -> None:
    from docx import Document
    from docx.enum.section import WD_ORIENT
    from docx.enum.text import WD_ALIGN_PARAGRAPH
    from docx.shared import Cm, Pt

    document = Document()
    section = document.sections[0]
    section.orientation = WD_ORIENT.LANDSCAPE
    section.page_width, section.page_height = section.page_height, section.page_width
    section.top_margin = Cm(1.4)
    section.bottom_margin = Cm(1.4)
    section.left_margin = Cm(1.2)
    section.right_margin = Cm(1.2)

    normal = document.styles["Normal"]
    normal.font.name = "Arial"
    normal.font.size = Pt(9)

    title = document.add_paragraph()
    title.alignment = WD_ALIGN_PARAGRAPH.CENTER
    title_run = title.add_run("THỐNG KÊ SỐ LƯỢNG TỆP GRIB2 CFS REFORECAST 90 NGÀY")
    title_run.bold = True
    title_run.font.name = "Arial"
    title_run.font.size = Pt(14)

    note = document.add_paragraph(
        "Phạm vi: 01/01/1999–31/03/2011. Mỗi ô là tổng số tệp có đường dẫn "
        "kết thúc bằng .grb2 trong toàn bộ các thư mục ngày của tháng tương ứng. "
        "Dấu 0 ở các tháng 4–12 năm 2011 thể hiện ngoài phạm vi dữ liệu của kho."
    )
    note.alignment = WD_ALIGN_PARAGRAPH.JUSTIFY

    add_table(document, "Bảng 1. Số lượng tệp GRIB2 của FLX", monthly["FLX"])
    add_table(document, "Bảng 2. Số lượng tệp GRIB2 của PGB", monthly["PGB"])

    source = document.add_paragraph()
    source.add_run("Nguồn: ").bold = True
    source.add_run(
        "NOAA/NCEI Climate Forecast System Reforecast; thống kê từ catalog THREDDS "
        "chính thức tương ứng với hai thư mục 6-hourly-flux-90-day-runs và "
        "6-hourly-by-pressure-level-90-day-runs. Ngày truy cập: 01/10/2026."
    )
    document.save(DOCX_FILE)


def print_summary(monthly: dict[str, dict[int, list[int]]]) -> None:
    for label in DATASETS:
        total = sum(sum(row) for row in monthly[label].values())
        daily_values = list(load_cache()[label].values())
        print(
            f"{label}: total={total:,}; daily min={min(daily_values):,}; "
            f"daily max={max(daily_values):,}"
        )
    print(f"CSV:  {CSV_FILE.name}")
    print(f"DOCX: {DOCX_FILE.name}")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--workers", type=int, default=32)
    args = parser.parse_args()
    cache = collect(max(1, min(args.workers, 64)))
    monthly = monthly_counts(cache)
    write_csv(monthly)
    write_docx(monthly)
    print_summary(monthly)


if __name__ == "__main__":
    main()