from __future__ import annotations

from ast import main
import multiprocessing
import shutil
from collections import defaultdict
from datetime import datetime, timedelta
from pathlib import Path

import pandas as pd
from log import logger

# base_dir = Path('/home/wzl/下载/cms')
base_dir = Path('/data/wind_files/headquarter/cms/')

tmp_dir = base_dir / 'tmp'
tmp_bak_dir = base_dir / 'tmp_bak'
days_dir_30 = base_dir / '30days'
days_dir_total = base_dir / 'total'

GROUP_COLUMNS = ['pointName', 'indexCode', 'departmentName']
REQUIRED_COLUMNS = [*GROUP_COLUMNS, 'localTime', 'value']


def get_files_by_relative_path() -> tuple[dict[Path, list[tuple[str, Path]]], str | None]:
    """Collect all YYYYMMDD folders under tmp, ordered chronologically per file."""
    files_by_path: dict[Path, list[tuple[str, Path]]] = defaultdict(list)
    date_dirs: list[tuple[datetime, Path]] = []
    logger.info(f'扫描 tmp 目录: {tmp_dir}')
    for date_dir in tmp_dir.iterdir() if tmp_dir.exists() else []:
        if not date_dir.is_dir():
            continue
        try:
            date_dirs.append((datetime.strptime(date_dir.name, '%Y%m%d'), date_dir))
        except ValueError:
            continue

    for _, date_dir in sorted(date_dirs):
        for file_path in date_dir.rglob('*.parquet'):
            files_by_path[file_path.relative_to(date_dir)].append((date_dir.name, file_path))

    for files in files_by_path.values():
        files.sort(key=lambda item: item[0])
    latest_date = max(date_dirs, key=lambda item: item[0])[1].name if date_dirs else None
    logger.info(f'扫描到 {len(files_by_path)} 个相对路径，latest_date={latest_date}')
    return files_by_path, latest_date


def read_parquet(file_path: Path) -> pd.DataFrame:
    df = pd.read_parquet(file_path)
    missing_columns = set(REQUIRED_COLUMNS) - set(df.columns)
    if missing_columns:
        raise ValueError(f'{file_path} 缺少字段: {sorted(missing_columns)}')
    return df[REQUIRED_COLUMNS].copy()


def complete_minute_data_fast(
    df: pd.DataFrame,
    query_date: str,
    history: pd.DataFrame | None = None,
) -> pd.DataFrame:
    if df.empty:
        return df.copy()

    query_day = pd.Timestamp(query_date)
    next_day = query_day + pd.Timedelta(days=1)
    df = df.copy()
    df['localTime'] = pd.to_datetime(df['localTime']).dt.floor('min')
    df = df.loc[(df['localTime'] >= query_day) & (df['localTime'] < next_day)]
    if df.empty:
        return df

    df = df.sort_values([*GROUP_COLUMNS, 'localTime'])
    df = df.drop_duplicates([*GROUP_COLUMNS, 'localTime'], keep='last')

    minute_range = pd.date_range(query_day, periods=24 * 60, freq='min')
    unique_keys = df[GROUP_COLUMNS].drop_duplicates().reset_index(drop=True)
    full_template = unique_keys.merge(
        pd.DataFrame({'localTime': minute_range}), how='cross'
    )
    completed = full_template.merge(
        df, on=[*GROUP_COLUMNS, 'localTime'], how='left'
    )

    # 用前一天最后一分钟的数据补全当天 00:00 的缺失值
    if history is not None and not history.empty:
        prior = history.copy()
        prior['localTime'] = pd.to_datetime(prior['localTime'])
        prior = prior.loc[prior['localTime'] < query_day, [*GROUP_COLUMNS, 'localTime', 'value']]
        if not prior.empty:
            prior = prior.sort_values([*GROUP_COLUMNS, 'localTime']).drop_duplicates(GROUP_COLUMNS, keep='last')
            seed = prior[GROUP_COLUMNS + ['value']].rename(columns={'value': 'seed_value'})
            completed = completed.merge(seed, on=GROUP_COLUMNS, how='left')
            first_minute = completed['localTime'].eq(query_day)
            needs_seed = first_minute & completed['value'].isna()
            completed['value'] = completed['value'].mask(needs_seed, completed['seed_value'])
            completed = completed.drop(columns=['seed_value'])
    else:
        completed['value'] = completed.groupby(GROUP_COLUMNS, dropna=False, sort=False)[    'value'].bfill()


    completed['value'] = completed.groupby(GROUP_COLUMNS, dropna=False, sort=False)['value'].ffill()
    completed['localTime'] = completed['localTime'].dt.strftime('%Y-%m-%d %H:%M:00')
    return completed


def remove_dates(df: pd.DataFrame, dates: set[str]) -> pd.DataFrame:
    if df.empty:
        return df
    local_dates = pd.to_datetime(df['localTime']).dt.strftime('%Y%m%d')
    return df.loc[~local_dates.isin(dates)].copy()


def combine_cms_data(
    relative_path: Path,
    files: list[tuple[str, Path]],
    latest_date: str,
    generate_total: bool = True,
) -> None:
    logger.info(f'{relative_path}: 开始处理，共 {len(files)} 天')
    total_path = days_dir_total / relative_path
    recent_path = days_dir_30 / relative_path
    existing_recent = read_parquet(recent_path) if recent_path.exists() else pd.DataFrame(columns=REQUIRED_COLUMNS)
    history = existing_recent.copy()
    completed_days: list[pd.DataFrame] = []

    for query_date, file_path in files:
        query_day = pd.Timestamp(datetime.strptime(query_date, '%Y%m%d'))
        completed = complete_minute_data_fast(read_parquet(file_path), query_day.strftime('%Y-%m-%d'), history)
        completed_days.append(completed)
        history = pd.concat([history, completed], ignore_index=True)
        bak_path = tmp_bak_dir / query_date / relative_path
        bak_path.parent.mkdir(parents=True, exist_ok=True)
        shutil.move(str(file_path), str(bak_path))
        logger.debug(f'{relative_path}: {query_date} 完成，已移动到 {bak_path}')

    input_dates = {date for date, _ in files}
    combined = pd.concat(
        [remove_dates(existing_recent, input_dates), *completed_days], ignore_index=True
    )
    combined['localTime'] = pd.to_datetime(combined['localTime'])
    combined = combined.sort_values([*GROUP_COLUMNS, 'localTime']).drop_duplicates(
        [*GROUP_COLUMNS, 'localTime'], keep='last'
    )
    combined['localTime'] = combined['localTime'].dt.strftime('%Y-%m-%d %H:%M:00')

    if generate_total:
        total_path.parent.mkdir(parents=True, exist_ok=True)
        combined.to_parquet(total_path, index=False)
        logger.info(f'{relative_path}: total 已写入，共 {len(combined)} 条')
    else:
        logger.info(f'{relative_path}: 跳过 total 写入')

    latest_day = pd.Timestamp(datetime.strptime(latest_date, '%Y%m%d'))
    recent_start = latest_day - timedelta(days=29)
    recent = combined.loc[pd.to_datetime(combined['localTime']) >= recent_start].copy()
    recent_path.parent.mkdir(parents=True, exist_ok=True)
    recent.to_parquet(recent_path, index=False)
    logger.info(f'{relative_path}: 完成，30days {len(recent)} 条')


def main(generate_total: bool = False):
    days_dir_total.mkdir(parents=True, exist_ok=True)
    days_dir_30.mkdir(parents=True, exist_ok=True)
    files_by_path, latest_date = get_files_by_relative_path()
    logger.info(f'待处理文件: {len(files_by_path)}，generate_total={generate_total}')
    if latest_date is None:
        raise SystemExit('tmp 下没有 YYYYMMDD 日期文件夹')
    with multiprocessing.Pool(processes=4, maxtasksperchild=5) as pool:
        pool.starmap(
            combine_cms_data,
            [
                (relative_path, files, latest_date, generate_total)
                for relative_path, files in files_by_path.items()
            ],
        )


if __name__ == '__main__':
    main(generate_total=True)