from __future__ import annotations

import os
import shutil
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Iterable, List, Optional, Tuple

from obspy import UTCDateTime, read  # type: ignore

# 根目录（递归查找波形文件）
ROOT = r"/mnt/seismic/ex_seismic/250815/102/102BHZ"

# 文件名后缀（既匹配 ...BHZ 也匹配 ....BHZ）
SUFFIX = ".BHZ"

# 输出根目录（按月分类的文件夹会创建在这里）
OUT_ROOT = os.path.join(ROOT, "_by_month")

# True=移动文件；False=复制文件
MOVE_FILES = True

# True=只打印不移动/复制
DRY_RUN = False


@dataclass(frozen=True)
class FileTimeRange:
    path: str
    start: UTCDateTime
    end: UTCDateTime


def _iter_files(root: str) -> Iterable[str]:
    root = os.path.abspath(root)
    for dirpath, _dirnames, filenames in os.walk(root):
        for fn in filenames:
            yield os.path.join(dirpath, fn)


def _match_suffix(path: str, suffix: str) -> bool:
    base = os.path.basename(path)
    return base.endswith(suffix) or base.endswith(f".{suffix}")


def _safe_datetime_utc(ts: float) -> datetime:
    return datetime.fromtimestamp(ts, tz=timezone.utc)


def _parse_time_range(path: str) -> Tuple[Optional[FileTimeRange], Optional[str]]:
    try:
        st = read(path, headonly=True, format="SAC")
    except Exception as e:
        return None, f"read(headonly) 失败: {e}"

    starts: List[UTCDateTime] = []
    ends: List[UTCDateTime] = []
    for tr in st:
        try:
            starts.append(tr.stats.starttime)
            ends.append(tr.stats.endtime)
        except Exception:
            continue

    if not starts or not ends:
        return None, "无有效 trace 的 start/end"

    return FileTimeRange(path=path, start=min(starts), end=max(ends)), None


def _month_key(t: UTCDateTime) -> str:
    dt = _safe_datetime_utc(t.timestamp)
    return dt.strftime("%Y%m")


def _ensure_dir(path: str) -> None:
    if DRY_RUN:
        return
    os.makedirs(path, exist_ok=True)


def _unique_dest(dest_dir: str, filename: str) -> str:
    base, ext = os.path.splitext(filename)
    candidate = os.path.join(dest_dir, filename)
    if not os.path.exists(candidate):
        return candidate
    i = 1
    while True:
        alt = os.path.join(dest_dir, f"{base}_{i}{ext}")
        if not os.path.exists(alt):
            return alt
        i += 1


def _move_or_copy(src: str, dest: str) -> None:
    if DRY_RUN:
        print(f"[DRY_RUN] {'MOVE' if MOVE_FILES else 'COPY'} {src} -> {dest}")
        return
    if MOVE_FILES:
        shutil.move(src, dest)
    else:
        shutil.copy2(src, dest)


def main() -> None:
    paths = [p for p in _iter_files(ROOT) if _match_suffix(p, SUFFIX)]
    paths.sort()
    if not paths:
        raise FileNotFoundError(f"在 {ROOT!r} 下没有找到以 {SUFFIX!r} 结尾的文件。")

    print(f"找到 {len(paths)} 个匹配文件，开始解析时间段并按月归类...")

    out_index = os.path.join(OUT_ROOT, "index.tsv")
    if not DRY_RUN:
        _ensure_dir(OUT_ROOT)

    moved = 0
    failed = 0
    error_lines: List[str] = []

    # index.tsv: path	start_utc	end_utc	month
    index_lines: List[str] = []

    for p in paths:
        info, err = _parse_time_range(p)
        if info is None:
            failed += 1
            if err:
                error_lines.append(f"{p}\t{err}")
            continue
        month = _month_key(info.start)
        month_dir = os.path.join(OUT_ROOT, month)
        _ensure_dir(month_dir)

        dest = _unique_dest(month_dir, os.path.basename(p))
        _move_or_copy(p, dest)
        moved += 1

        start_dt = _safe_datetime_utc(info.start.timestamp).isoformat()
        end_dt = _safe_datetime_utc(info.end.timestamp).isoformat()
        index_lines.append(f"{p}\t{start_dt}\t{end_dt}\t{month}")

    if index_lines and not DRY_RUN:
        with open(out_index, "w", encoding="utf-8") as f:
            f.write("path\tstart_utc\tend_utc\tmonth\n")
            for line in index_lines:
                f.write(line + "\n")

    print(f"完成：成功处理 {moved} 个文件，失败 {failed} 个。")
    if not DRY_RUN:
        print(f"索引写入：{out_index}")
        if error_lines:
            err_path = os.path.join(OUT_ROOT, "errors.tsv")
            with open(err_path, "w", encoding="utf-8") as f:
                f.write("path\treason\n")
                for line in error_lines:
                    f.write(line + "\n")
            print(f"错误明细：{err_path}")


if __name__ == "__main__":
    main()