def find_partition_files(input_dir: str) -> list[str]:
    """Find all partition metric files in the input directory."""
    pattern = os.path.join(input_dir, "**/metrics-*.json")
    return glob.glob(pattern, recursive=True)

def merge_metrics(
    input_dir: str,
    output_file: str,
    run_id: str,
    commit_sha: str,
    branch: str | None = None,
) -> bool:
    """Merge all partition metrics into a consolidated file."""
    run_date = datetime.now(timezone.utc).isoformat()

    # Find all partition files
    partition_files = find_partition_files(input_dir)
    print(f"Found {len(partition_files)} partition file(s)")

    all_results = []
    if not partition_files:
        print("No partition metrics files found")
    else:
        # Load all partition files
        for filepath in sorted(partition_files):
            print(f"  Reading: {filepath}")
            metrics = load_partition_metrics(filepath)
            if metrics and "results" in metrics:
                all_results.extend(metrics["results"])
        print(f"Total results collected: {len(all_results)}")

    # Create consolidated structure
    consolidated = {
        "run_id": run_id,
        "run_date": run_date,
        "commit_sha": commit_sha,
        "branch": branch,
        "results": all_results,
    }

    # Ensure output directory exists and write output
    try:
        os.makedirs(os.path.dirname(output_file) or ".", exist_ok=True)
        with open(output_file, "w", encoding="utf-8") as f:
            json.dump(consolidated, f, indent=2)

        if not partition_files:
            print(f"Created empty consolidated file: {output_file}")
        else:
            print(f"Saved consolidated metrics to: {output_file}")
        return True
    except OSError as e:
        print(f"Error writing consolidated file: {e}")
        return False

def find_result_files(search_dirs: list[str]) -> list[str]:
    """Find all results_*.json files in the given directories."""
    result_files = set()
    for search_dir in search_dirs:
        if os.path.exists(search_dir):
            pattern = os.path.join(search_dir, "**/results_*.json")
            result_files.update(glob.glob(pattern, recursive=True))
    return list(result_files)

        "latency_ms": latency * 1000 if latency is not None else None,
        "input_throughput": result.get("input_throughput"),
        "output_throughput": result.get("output_throughput"),
        "overall_throughput": result.get("overall_throughput"),
        "ttft_ms": last_ttft * 1000 if last_ttft is not None else None,

def save_metrics(
    gpu_config: str,
    partition: int,
    run_id: str,
    output_file: str,
    search_dirs: list[str],
) -> bool:
    """Collect metrics and save to output file."""
    timestamp = datetime.now(timezone.utc).isoformat()

    # Find all result files
    result_files = find_result_files(search_dirs)
    print(f"Found {len(result_files)} result file(s)")

    grouped = []
    if not result_files:
        print("No benchmark result files found")
    else:
        # Parse all result files
        all_results = []
        for filepath in sorted(result_files):
            print(f"  Reading: {filepath}")
            results = parse_result_file(filepath)
            all_results.extend(results)
        print(f"Total benchmark results: {len(all_results)}")

        # Group by model/variant
        grouped = group_results_by_model(all_results, gpu_config, partition)

    # Create metrics structure
    metrics = {
        "run_id": run_id,
        "timestamp": timestamp,
        "gpu_config": gpu_config,
        "partition": partition,
        "results": grouped,
    }

    # Ensure output directory exists and write output
    try:
        os.makedirs(os.path.dirname(output_file) or ".", exist_ok=True)
        with open(output_file, "w", encoding="utf-8") as f:
            json.dump(metrics, f, indent=2)

        if not result_files:
            print(f"Created empty metrics file: {output_file}")
        else:
            print(f"Saved metrics to: {output_file}")
        return True
    except OSError as e:
        print(f"Error writing metrics file: {e}")
        return False