import argparse
import re
from pathlib import Path
import xml.etree.ElementTree as ET
from typing import Dict, Tuple, Optional


def indent(elem: ET.Element, level: int = 0) -> None:
    i = "\n" + level * "    "
    if len(elem):
        if not elem.text or not elem.text.strip():
            elem.text = i + "    "
        for e in elem:
            indent(e, level + 1)
        if not e.tail or not e.tail.strip():
            e.tail = i
    if level and (not elem.tail or not elem.tail.strip()):
        elem.tail = i


_KV_MA_RE = re.compile(r"(?P<kv>\d+)\s*[kK][vV]_(?P<ma>\d+)\s*[mM][aA]")


def _parse_folder_name(folder_name: str) -> Tuple[Optional[int], Optional[int]]:
    m = _KV_MA_RE.search(folder_name)
    if not m:
        return None, None
    try:
        return int(m.group("kv")), int(m.group("ma"))
    except Exception:
        return None, None


def collect_csvs(csv_dir: Path) -> Dict[str, str]:
    """Return dict with keys 'results', 'lung', 'advanced' mapping to file paths or ''."""
    out = {"results": "", "lung": "", "advanced": ""}
    if not csv_dir.exists() or not csv_dir.is_dir():
        return out
    for p in sorted(csv_dir.iterdir()):
        if not p.is_file():
            continue
        name = p.name.lower()
        if name.endswith("analysis_results.csv") and "lung" not in name:
            out["results"] = str(p.resolve())
        elif name.endswith("lung_results.csv"):
            out["lung"] = str(p.resolve())
        elif name.endswith("advanced_metrics.csv"):
            out["advanced"] = str(p.resolve())
    return out


def main(src_dir: Path, out_file: Path) -> None:
    src_dir = src_dir.expanduser().resolve()
    out_file = out_file.expanduser().resolve()

    if not src_dir.exists() or not src_dir.is_dir():
        raise SystemExit(f"Source directory does not exist or is not a directory: {src_dir}")

    root = ET.Element("experiments")

    entries = []
    for folder in sorted(src_dir.iterdir()):
        if not folder.is_dir():
            continue
        csv_dir = folder / "csv"
        png_dir = folder / "png"
        csvs = collect_csvs(csv_dir)

        kv, ma = _parse_folder_name(folder.name)
        entries.append(
            {
                "folder": folder,
                "name": folder.name,
                "kv": str(kv) if kv is not None else "",
                "ma": str(ma) if ma is not None else "",
                "results": csvs["results"],
                "lung": csvs["lung"],
                "advanced": csvs["advanced"],
                "png_dir": str(png_dir.resolve()) if png_dir.exists() else "",
            }
        )

    # sort: numeric kv then numeric ma if present, fallback to folder name
    def _sort_key(e):
        try:
            return (int(e["kv"]) if e["kv"] else 9999, int(e["ma"]) if e["ma"] else 9999, e["name"])
        except Exception:
            return (9999, 9999, e["name"])

    for e in sorted(entries, key=_sort_key):
        attrib = {
            "name": e["name"],
            "path": e["results"],
            "lung_path": e["lung"],
            "advanced_path": e["advanced"],
            "status_plot": "enhanced",
        }
        # remove empty attributes to keep XML clean
        attrib_filtered: dict[str, str] = {k: str(v) for k, v in attrib.items() if v}
        ET.SubElement(root, "experiment", attrib_filtered)

    indent(root)
    tree = ET.ElementTree(root)
    tree.write(out_file, encoding="utf-8", xml_declaration=True)
    print(f"Wrote {len(list(root))} experiments to {out_file}")


if __name__ == "__main__":
    p = argparse.ArgumentParser(
        description="Generate experiments XML from a directory tree with subfolders like '100kV_100mA'"
    )
    p.add_argument("src_dir", type=Path)
    p.add_argument("out_file", type=Path, nargs="?", default=Path("experiments.xml"))
    args = p.parse_args()
    main(args.src_dir, args.out_file)
