import json
import re
import numpy as np
import open3d as o3d
from pathlib import Path
from .common import Config, OptimizerContext

def run_module0_metadata(cfg: Config, ctx: OptimizerContext):
    print("\n=== [Module 0] Load Metadata & Map Files ===")
    if not cfg.gnss_json.exists():
        raise FileNotFoundError(f"JSON not found: {cfg.gnss_json}")

    with open(cfg.gnss_json, "r", encoding="utf-8") as f:
        ctx.gnss_info = json.load(f)

    def get_idx(k: str) -> int:
        try:
            return int(k.split("_")[-1])
        except Exception:
            return 999999

    ctx.block_ids = sorted(ctx.gnss_info.keys(), key=get_idx)
    print(f"[Info] Blocks found: {ctx.block_ids}")

    # Map PCD files
    pcd_files = list(cfg.pcd_dir.glob("*.pcd"))
    if len(pcd_files) == 0:
        raise FileNotFoundError(f"No pcd found in {cfg.pcd_dir}")

    mapped_count = 0
    for bid in ctx.block_ids:
        idx = get_idx(bid)
        found = None
        for p in pcd_files:
            m = re.search(r"(\d+)\.pcd$", p.name, re.IGNORECASE)
            if m and int(m.group(1)) == idx:
                found = p
                break
        
        if found:
            ctx.pcd_paths[bid] = found
            mapped_count += 1
            print(f"[Map] {bid} -> PCD {found.name}")
        else:
            print(f"[Warn] No PCD found for {bid}")

    if mapped_count == 0:
        raise RuntimeError("No PCD files matched to blocks!")

    # Map Trajectories
    for bid in ctx.block_ids:
        info = ctx.gnss_info[bid]
        if "pose_file" not in info or "gnss_file" not in info:
            raise RuntimeError(f"{bid} missing pose_file/gnss_file in json.")
        
        ctx.pose_paths[bid] = cfg.pose_root_dir / info["pose_file"]
        ctx.gnss_kf_paths[bid] = cfg.pose_root_dir / info["gnss_file"]
        print(f"[Map] {bid} -> pose {ctx.pose_paths[bid].name}, gnss {ctx.gnss_kf_paths[bid].name}")

def run_module1_load_pcd(cfg: Config, ctx: OptimizerContext):
    print("\n=== [Module 1] Load Preprocessed PCDs ===")
    for bid in ctx.block_ids:
        if bid not in ctx.pcd_paths:
            continue
        path = ctx.pcd_paths[bid]
        pcd = o3d.io.read_point_cloud(str(path))
        n = np.asarray(pcd.points).shape[0]
        print(f"[Load] {bid}: {path.name} points={n}")
        ctx.pcd_ds[bid] = pcd
