#!/usr/bin/env python
import argparse
import os
from pathlib import Path

import pandas as pd
import trimesh


def parse_args():
    parser = argparse.ArgumentParser(
        description="Extract basic geometric features from STL file(s) using trimesh."
    )
    parser.add_argument(
        "--stl_path",
        help="Path to a single STL file or a folder containing STL files (non-recursive). Defaults to current directory.",
        default=".",
    )
    parser.add_argument(
        "--output_path",
        help="Folder where the resulting CSV will be saved. Defaults to current directory.",
        default=".",
    )
    return parser.parse_args()


def collect_stl_files(stl_path: Path):
    if stl_path.is_file():
        if stl_path.suffix.lower() != ".stl":
            raise ValueError(f"Provided file is not an STL: {stl_path}")
        return [stl_path]

    if stl_path.is_dir():
        files = []
        for name in os.listdir(stl_path):
            p = stl_path / name
            if p.is_file() and p.suffix.lower() == ".stl":
                files.append(p)
        if not files:
            raise ValueError(f"No .stl files found in folder: {stl_path}")
        return files

    raise FileNotFoundError(f"Path does not exist: {stl_path}")


def extract_features(stl_file: Path):
    mesh = trimesh.load(stl_file, force="mesh")
    if mesh.is_empty:
        raise ValueError(f"Empty mesh in file: {stl_file}")

    file_weight_kb = os.path.getsize(stl_file) / 1024.0
    bbox_x, bbox_y, bbox_z = mesh.bounding_box.extents

    volume = mesh.volume
    surface_area = mesh.area

    stl = {
        "file_name": stl_file.name,
        "file_path": str(stl_file),
        "kb": file_weight_kb,
        "volume": volume,
        "surface_area": surface_area,
        "scale": mesh.scale,
        "bbox_area": bbox_x * bbox_y * bbox_z,
        "euler_number": mesh.euler_number,
        "surface_volume_ratio": surface_area / volume if volume != 0 else None,
    }
    return stl


def main():
    args = parse_args()

    stl_path = Path(args.stl_path).expanduser().resolve()
    output_dir = Path(args.output_path).expanduser().resolve()

    output_dir.mkdir(parents=True, exist_ok=True)

    stl_files = collect_stl_files(stl_path)
    total = len(stl_files)

    rows = []
    for i, f in enumerate(stl_files, start=1):
        try:
            features = extract_features(f)
            rows.append(features)
        except Exception as e:
            print(f"\n[WARN] Skipping {f}: {e}")
        # simple progress bar
        pct = i / total * 100
        bar_len = 30
        filled = int(bar_len * i / total)
        bar = "#" * filled + "-" * (bar_len - filled)
        print(f"\r[{bar}] {i}/{total} ({pct:5.1f}%)", end="", flush=True)

    print()  # newline after progress bar

    if not rows:
        raise RuntimeError("No valid STL files processed; nothing to save.")

    df = pd.DataFrame(rows)

    out_csv = output_dir / "stl_features.csv"
    df.to_csv(out_csv, index=False)
    print(f"Saved {len(df)} rows to {out_csv}")


if __name__ == "__main__":
    main()
