#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import json
from pathlib import Path
from collections import Counter, defaultdict
import argparse


def analisar_dataset(root: Path):
    resultado = {}

    for grupo_dir in sorted([p for p in root.iterdir() if p.is_dir()]):
        labels_dir = grupo_dir / "labels"
        if not labels_dir.is_dir():
            continue

        contagem = Counter()
        total = 0

        for json_path in sorted(labels_dir.glob("*.json")):
            try:
                with json_path.open("r", encoding="utf-8") as f:
                    data = json.load(f)

                status = data.get("estado_corredor")
                if not status:
                    status = "SEM_STATUS"

                contagem[status] += 1
                total += 1

            except Exception:
                contagem["ERRO_JSON"] += 1
                total += 1

        resultado[grupo_dir.name] = {
            "total": total,
            "contagem": contagem
        }

    return resultado


def imprimir_relatorio(resultado):
    todos_status = sorted({
        status
        for dados in resultado.values()
        for status in dados["contagem"].keys()
    })

    print("\nRELATÓRIO DE STATUS DO CORREDOR")
    print("=" * 80)

    total_geral = 0
    contagem_geral = Counter()

    for grupo, dados in resultado.items():
        total = dados["total"]
        contagem = dados["contagem"]

        total_geral += total
        contagem_geral.update(contagem)

        print(f"\nGrupo: {grupo}")
        print(f"Total labels: {total}")

        if total == 0:
            print("  Nenhum label encontrado.")
            continue

        for status in todos_status:
            qtd = contagem.get(status, 0)
            perc = (qtd / total) * 100
            print(f"  {status:<20} {qtd:>5}  {perc:>6.2f}%")

    print("\nGERAL")
    print("=" * 80)
    print(f"Total labels: {total_geral}")

    if total_geral > 0:
        for status in todos_status:
            qtd = contagem_geral.get(status, 0)
            perc = (qtd / total_geral) * 100
            print(f"  {status:<20} {qtd:>5}  {perc:>6.2f}%")


def salvar_csv(resultado, out_csv: Path):
    import csv

    todos_status = sorted({
        status
        for dados in resultado.values()
        for status in dados["contagem"].keys()
    })

    with out_csv.open("w", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        header = ["grupo", "total"]

        for status in todos_status:
            header.append(f"{status}_qtd")
            header.append(f"{status}_perc")

        writer.writerow(header)

        for grupo, dados in resultado.items():
            total = dados["total"]
            contagem = dados["contagem"]

            row = [grupo, total]
            for status in todos_status:
                qtd = contagem.get(status, 0)
                perc = (qtd / total) * 100 if total > 0 else 0
                row.extend([qtd, round(perc, 2)])

            writer.writerow(row)

    print(f"\nCSV salvo em: {out_csv}")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--root",
        required=True,
        help="Pasta original/group contendo os grupos com labels/"
    )
    parser.add_argument(
        "--csv",
        default=None,
        help="Opcional: caminho para salvar o relatório CSV"
    )

    args = parser.parse_args()

    root = Path(args.root)
    resultado = analisar_dataset(root)

    imprimir_relatorio(resultado)

    if args.csv:
        salvar_csv(resultado, Path(args.csv))


if __name__ == "__main__":
    main()