#!/usr/bin/env python3
"""
Audita e organiza as amostras multiespectrais brutas.

Regras por amostra:
  - completa/mapeada: CAM_A + CAM_B + CAM_C + meta JSON + preview PNG + mask PNG;
  - completa/pendente: CAM_A + CAM_B + CAM_C + meta JSON + preview PNG, sem mask;
  - incompleta: falta qualquer item obrigatório acima (exceto a mask).

Por padrão o programa apenas simula. Com --aplicar, copia os previews pendentes
e move conjuntos incompletos para uma quarentena recuperável. A exclusão
definitiva exige também --excluir-definitivamente.
"""

from __future__ import annotations

import argparse
import csv
import hashlib
import re
import shutil
import sys
from collections import defaultdict
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Iterable


GRUPOS_PADRAO = (
    "cana",
    "cana_erva",
    "chao",
    "chao_cana",
    "chao_cana_erva",
    "chao_erva",
    "erva",
)

SUBPASTAS = ("bins", "masks", "metas", "previews")
RE_COPIA_WINDOWS = re.compile(r"\s*\((\d+)\)\s*$")
RE_CAM = re.compile(r"(?i)(?:^|[_\-\s])CAM[_\-\s]*([ABC])(?:$|[_\-\s])")


@dataclass(frozen=True, order=True)
class ChaveAmostra:
    base: str
    copia: int

    @property
    def texto(self) -> str:
        return self.base if self.copia == 0 else f"{self.base} ({self.copia})"


@dataclass
class Amostra:
    grupo: str
    chave: ChaveAmostra
    cams: dict[str, list[Path]] = field(
        default_factory=lambda: {"A": [], "B": [], "C": []}
    )
    metas: list[Path] = field(default_factory=list)
    previews: list[Path] = field(default_factory=list)
    masks: list[Path] = field(default_factory=list)
    outros: list[Path] = field(default_factory=list)

    @property
    def arquivos(self) -> list[Path]:
        resultado: list[Path] = []
        for camera in "ABC":
            resultado.extend(self.cams[camera])
        resultado.extend(self.metas)
        resultado.extend(self.previews)
        resultado.extend(self.masks)
        resultado.extend(self.outros)
        return resultado

    @property
    def conflitos(self) -> list[str]:
        erros: list[str] = []
        for camera in "ABC":
            if len(self.cams[camera]) > 1:
                erros.append(f"{len(self.cams[camera])} arquivos CAM_{camera}")
        if len(self.metas) > 1:
            erros.append(f"{len(self.metas)} metas")
        if len(self.previews) > 1:
            erros.append(f"{len(self.previews)} previews")
        if len(self.masks) > 1:
            erros.append(f"{len(self.masks)} masks")
        return erros

    @property
    def faltando(self) -> list[str]:
        faltas = [f"CAM_{c}" for c in "ABC" if len(self.cams[c]) == 0]
        if not self.metas:
            faltas.append("meta")
        if not self.previews:
            faltas.append("preview")
        return faltas

    @property
    def estado(self) -> str:
        if self.conflitos:
            return "CONFLITO"
        if self.faltando:
            return "INCOMPLETA"
        if self.masks:
            return "PRONTA"
        return "PENDENTE_MASK"


def separar_numero_copia(stem: str) -> tuple[str, int]:
    match = RE_COPIA_WINDOWS.search(stem)
    if not match:
        return stem.strip(), 0
    return stem[: match.start()].strip(), int(match.group(1))


def limpar_base(stem: str, tipo: str) -> ChaveAmostra:
    sem_copia, copia = separar_numero_copia(stem)

    # Alguns pipelines incluem CAM_A também no nome do preview/mask/meta.
    # Removê-lo em todos os tipos permite associar essas variações ao mesmo trio.
    sem_copia = RE_CAM.sub("_", sem_copia)

    if tipo == "meta":
        sem_copia = re.sub(r"(?i)(?:[_\-\s]+)(?:meta|metadata)$", "", sem_copia)
    elif tipo == "preview":
        sem_copia = re.sub(r"(?i)(?:[_\-\s]+)preview$", "", sem_copia)
    elif tipo == "mask":
        # A ferramenta de anotação pode devolver a mask com o próprio nome do
        # preview copiado. Portanto, em masks aceitamos tanto _mask quanto
        # _preview como sufixo da mesma amostra.
        sem_copia = re.sub(r"(?i)(?:[_\-\s]+)(?:mask|preview)$", "", sem_copia)

    sem_copia = re.sub(r"[_\-\s]+$", "", sem_copia)
    sem_copia = re.sub(r"[_\-\s]+", "_", sem_copia).strip("_").casefold()
    return ChaveAmostra(sem_copia, copia)


def camera_do_bin(path: Path) -> str | None:
    sem_copia, _ = separar_numero_copia(path.stem)
    match = RE_CAM.search(sem_copia)
    return match.group(1).upper() if match else None


def arquivos_validos(pasta: Path, extensoes: set[str]) -> Iterable[Path]:
    if not pasta.exists():
        return ()
    return (
        p
        for p in pasta.iterdir()
        if p.is_file() and p.suffix.casefold() in extensoes
    )


def inventariar_grupo(raiz: Path, grupo: str) -> dict[ChaveAmostra, Amostra]:
    pasta_grupo = raiz / grupo
    amostras: dict[ChaveAmostra, Amostra] = {}

    def obter(chave: ChaveAmostra) -> Amostra:
        return amostras.setdefault(chave, Amostra(grupo, chave))

    for path in arquivos_validos(pasta_grupo / "bins", {".bin"}):
        camera = camera_do_bin(path)
        chave = limpar_base(path.stem, "bin")
        amostra = obter(chave)
        if camera:
            amostra.cams[camera].append(path)
        else:
            amostra.outros.append(path)

    for path in arquivos_validos(pasta_grupo / "metas", {".json"}):
        obter(limpar_base(path.stem, "meta")).metas.append(path)

    for path in arquivos_validos(pasta_grupo / "previews", {".png"}):
        obter(limpar_base(path.stem, "preview")).previews.append(path)

    for path in arquivos_validos(pasta_grupo / "masks", {".png"}):
        obter(limpar_base(path.stem, "mask")).masks.append(path)

    return amostras


def hash_arquivo(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as arquivo:
        for bloco in iter(lambda: arquivo.read(1024 * 1024), b""):
            digest.update(bloco)
    return digest.hexdigest()


def destino_sem_sobrescrever(origem: Path, pasta_destino: Path) -> tuple[Path, str]:
    destino = pasta_destino / origem.name
    if not destino.exists():
        return destino, "copiar"
    if origem.stat().st_size == destino.stat().st_size:
        if hash_arquivo(origem) == hash_arquivo(destino):
            return destino, "ja_existe_igual"
    contador = 2
    while True:
        alternativo = pasta_destino / f"{origem.stem}__duplicado_{contador}{origem.suffix}"
        if not alternativo.exists():
            return alternativo, "copiar_renomeado"
        contador += 1


def mover_preservando_estrutura(path: Path, raiz: Path, quarentena: Path) -> Path:
    relativo = path.relative_to(raiz)
    destino = quarentena / relativo
    destino.parent.mkdir(parents=True, exist_ok=True)
    if destino.exists():
        destino = destino.with_name(
            f"{destino.stem}__conflito_{datetime.now():%H%M%S%f}{destino.suffix}"
        )
    shutil.move(str(path), str(destino))
    return destino


def escrever_relatorio(path: Path, linhas: list[dict[str, str]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    campos = [
        "grupo",
        "amostra",
        "estado",
        "faltando",
        "conflitos",
        "cams",
        "meta",
        "preview",
        "mask",
        "acao",
    ]
    with path.open("w", newline="", encoding="utf-8-sig") as arquivo:
        writer = csv.DictWriter(arquivo, fieldnames=campos, delimiter=";")
        writer.writeheader()
        writer.writerows(linhas)


def validar_raiz(raiz: Path, grupos: tuple[str, ...]) -> None:
    if not raiz.is_dir():
        raise SystemExit(f"ERRO: pasta de entrada não encontrada: {raiz}")
    ausentes = [grupo for grupo in grupos if not (raiz / grupo).is_dir()]
    if ausentes:
        raise SystemExit(
            "ERRO: faltam as pastas de grupo: " + ", ".join(ausentes)
        )
    for grupo in grupos:
        faltantes = [s for s in SUBPASTAS if not (raiz / grupo / s).is_dir()]
        if faltantes:
            raise SystemExit(
                f"ERRO: {raiz / grupo} não contém: {', '.join(faltantes)}"
            )


def executar(args: argparse.Namespace) -> int:
    raiz = args.entrada.resolve()
    saida = args.saida.resolve()
    grupos = tuple(args.grupos)
    validar_raiz(raiz, grupos)

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    quarentena = raiz.parent / f"_quarentena_incompletos_{timestamp}"
    linhas: list[dict[str, str]] = []
    totais: defaultdict[str, int] = defaultdict(int)

    print(f"Modo: {'APLICAR' if args.aplicar else 'SIMULAÇÃO (nenhum arquivo será alterado)'}")
    print(f"Entrada: {raiz}")
    print(f"Previews pendentes: {saida}")
    if args.aplicar and not args.excluir_definitivamente:
        print(f"Quarentena: {quarentena}")
    print()

    for grupo in grupos:
        amostras = inventariar_grupo(raiz, grupo)
        print(f"[{grupo}] {len(amostras)} amostras encontradas")

        for chave in sorted(amostras):
            amostra = amostras[chave]
            estado = amostra.estado
            totais[estado] += 1
            acao = "ignorar"

            if estado == "PENDENTE_MASK":
                destino, modo_copia = destino_sem_sobrescrever(
                    amostra.previews[0], saida / grupo
                )
                if modo_copia == "ja_existe_igual":
                    acao = f"preview já existe: {destino}"
                    totais["PREVIEW_JA_EXISTIA"] += 1
                else:
                    acao = f"copiar preview -> {destino}"
                    totais["PREVIEWS_A_COPIAR"] += 1
                    if args.aplicar:
                        destino.parent.mkdir(parents=True, exist_ok=True)
                        shutil.copy2(amostra.previews[0], destino)
                        totais["PREVIEWS_COPIADOS"] += 1

            elif estado == "INCOMPLETA":
                if args.excluir_definitivamente:
                    acao = f"excluir {len(amostra.arquivos)} arquivo(s) definitivamente"
                else:
                    acao = f"mover {len(amostra.arquivos)} arquivo(s) para quarentena"
                totais["ARQUIVOS_INCOMPLETOS"] += len(amostra.arquivos)
                if args.aplicar:
                    for arquivo in amostra.arquivos:
                        if args.excluir_definitivamente:
                            arquivo.unlink()
                        else:
                            mover_preservando_estrutura(arquivo, raiz, quarentena)
                    totais["AMOSTRAS_INCOMPLETAS_TRATADAS"] += 1

            elif estado == "CONFLITO":
                acao = "revisar manualmente; nada será removido"

            linhas.append(
                {
                    "grupo": grupo,
                    "amostra": chave.texto,
                    "estado": estado,
                    "faltando": ", ".join(amostra.faltando),
                    "conflitos": ", ".join(amostra.conflitos),
                    "cams": ", ".join(
                        f"{c}:{len(amostra.cams[c])}" for c in "ABC"
                    ),
                    "meta": str(len(amostra.metas)),
                    "preview": str(len(amostra.previews)),
                    "mask": str(len(amostra.masks)),
                    "acao": acao,
                }
            )

    relatorio = (
        args.relatorio.resolve()
        if args.relatorio
        else Path.cwd() / f"relatorio_limpeza_{timestamp}.csv"
    )
    escrever_relatorio(relatorio, linhas)

    print("\nResumo")
    for nome in (
        "PRONTA",
        "PENDENTE_MASK",
        "INCOMPLETA",
        "CONFLITO",
        "PREVIEWS_A_COPIAR",
        "PREVIEW_JA_EXISTIA",
        "PREVIEWS_COPIADOS",
        "ARQUIVOS_INCOMPLETOS",
        "AMOSTRAS_INCOMPLETAS_TRATADAS",
    ):
        print(f"  {nome}: {totais[nome]}")
    print(f"\nRelatório: {relatorio}")

    if not args.aplicar:
        print("\nSIMULAÇÃO concluída. Revise o CSV e rode novamente com --aplicar.")
    elif args.excluir_definitivamente:
        print("\nAplicação concluída com exclusão definitiva dos incompletos.")
    else:
        print("\nAplicação concluída. Os incompletos foram preservados na quarentena.")

    return 0


def criar_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Audita brutas multiespectrais e separa previews sem mask."
    )
    parser.add_argument(
        "--entrada",
        type=Path,
        default=Path("dataset/brutas/group"),
        help="Pasta que contém os sete grupos.",
    )
    parser.add_argument(
        "--saida",
        type=Path,
        default=Path(r"C:\Users\USER\Desktop\fotos_multiespectrais\mapear\group"),
        help="Destino dos previews que ainda não possuem mask.",
    )
    parser.add_argument(
        "--grupos",
        nargs="+",
        default=list(GRUPOS_PADRAO),
        help="Grupos que serão processados.",
    )
    parser.add_argument(
        "--relatorio",
        type=Path,
        help="Caminho opcional do relatório CSV.",
    )
    parser.add_argument(
        "--aplicar",
        action="store_true",
        help="Executa cópias e limpeza. Sem esta flag, apenas simula.",
    )
    parser.add_argument(
        "--excluir-definitivamente",
        action="store_true",
        help="Com --aplicar, apaga incompletos em vez de movê-los à quarentena.",
    )
    return parser


def main() -> int:
    parser = criar_parser()
    args = parser.parse_args()
    if args.excluir_definitivamente and not args.aplicar:
        parser.error("--excluir-definitivamente só pode ser usado com --aplicar")
    return executar(args)


if __name__ == "__main__":
    try:
        sys.exit(main())
    except KeyboardInterrupt:
        print("\nCancelado pelo usuário.", file=sys.stderr)
        sys.exit(130)