"""
============================================================================
 AIRCRAFT ROUTING - ESQUADRAO ARARA - V2
============================================================================
 Primeira versao do modelo OAMRP usando dados reais pre-processados do
 registro de voo 2025.

 Diferencas em relacao ao v1:
   - le o CSV consolidado gerado a partir da Planilha 2025;
   - filtra uma janela de datas, com janeiro/2025 como padrao;
   - usa aeronaves reais observadas no registro: 2800, 2803, 2809 e 2811;
   - gera um CSV de resultado em db/processed/.

 Esta versao ainda e uma rodada exploratoria. Todas as missoes recebem
 prioridade 1 e a base de manutencao considerada e SBMN.
============================================================================
"""

from __future__ import annotations

import argparse
import csv
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path

import pulp


BASE_DIR = Path(__file__).resolve().parents[1]
CSV_REGISTRO_PADRAO = BASE_DIR / "db" / "processed" / "registro_voo_2025_consolidado.csv"
CSV_RESULTADO_PADRAO = BASE_DIR / "db" / "processed" / "resultado_oamrp_jan_2025.csv"

DATA_INICIO_PADRAO = "2025-01-01"
DATA_FIM_PADRAO = "2025-02-01"  # exclusivo

BASE_MANUTENCAO = "SBMN"
BASE_INICIAL = BASE_MANUTENCAO
TURNAROUND_MIN_H = 1.0
LIMITE_HORAS_F = 300.0
PRIORIDADE_PADRAO = 1

AERONAVES = [
    # id    base inicial  horas acumuladas desde ultima inspecao
    ("2800", BASE_INICIAL, 0.0),
    ("2803", BASE_INICIAL, 0.0),
    ("2809", BASE_INICIAL, 0.0),
    ("2811", BASE_INICIAL, 0.0),
]


@dataclass(frozen=True)
class Missao:
    id: str
    orig: str
    dest: str
    partida_utc: datetime
    chegada_utc: datetime
    prioridade: int
    periodo: str
    linhas_origem: str
    aeronave_real: str
    om: str

    @property
    def dur_h(self) -> float:
        return (self.chegada_utc - self.partida_utc).total_seconds() / 3600


def parse_utc(valor: str) -> datetime:
    texto = str(valor).strip()
    if texto.endswith("Z"):
        texto = texto[:-1] + "+00:00"
    dt = datetime.fromisoformat(texto)
    if dt.tzinfo is None:
        dt = dt.replace(tzinfo=timezone.utc)
    return dt.astimezone(timezone.utc)


def parse_data_inicio(valor: str) -> datetime:
    return datetime.fromisoformat(valor).replace(tzinfo=timezone.utc)


def carregar_missoes(caminho: Path, inicio: datetime, fim: datetime) -> list[Missao]:
    missoes: list[Missao] = []
    with caminho.open(newline="", encoding="utf-8-sig") as arquivo:
        leitor = csv.DictReader(arquivo)
        for linha in leitor:
            partida = parse_utc(linha["partida_utc"])
            chegada = parse_utc(linha["chegada_utc"])
            if not (inicio <= partida < fim):
                continue
            if chegada <= partida:
                raise ValueError(f"Missao {linha['id_registro']}: chegada antes da partida")
            missoes.append(
                Missao(
                    id=linha["id_registro"],
                    orig=linha["localidade_dep"],
                    dest=linha["localidade_arr"],
                    partida_utc=partida,
                    chegada_utc=chegada,
                    prioridade=PRIORIDADE_PADRAO,
                    periodo=linha.get("periodo", ""),
                    linhas_origem=linha.get("linhas_origem", ""),
                    aeronave_real=linha.get("aeronave", ""),
                    om=linha.get("om", ""),
                )
            )
    return sorted(missoes, key=lambda m: (m.partida_utc, m.id))


def horas_relativas(missoes: list[Missao]) -> tuple[dict[str, dict[str, object]], datetime]:
    horizonte_inicio = min(m.partida_utc for m in missoes)
    dados = {}
    for missao in missoes:
        dados[missao.id] = {
            "orig": missao.orig,
            "dest": missao.dest,
            "part": (missao.partida_utc - horizonte_inicio).total_seconds() / 3600,
            "cheg": (missao.chegada_utc - horizonte_inicio).total_seconds() / 3600,
            "prio": missao.prioridade,
            "dur": missao.dur_h,
            "obj": missao,
        }
    return dados, horizonte_inicio


def resolver(missoes: list[Missao]) -> dict[str, object]:
    if not missoes:
        raise ValueError("Nenhuma missao encontrada na janela informada")

    miss, horizonte_inicio = horas_relativas(missoes)
    acft = {a[0]: {"base": a[1], "f": a[2]} for a in AERONAVES}
    ids = list(miss.keys())
    ks = list(acft.keys())

    conex = [
        (i, j)
        for i in ids
        for j in ids
        if i != j
        and miss[i]["dest"] == miss[j]["orig"]
        and miss[j]["part"] >= miss[i]["cheg"] + TURNAROUND_MIN_H
    ]

    def pode_iniciar(k: str, i: str) -> bool:
        return miss[i]["orig"] == acft[k]["base"]

    mdl = pulp.LpProblem("ARARA_OAMRP_v2", pulp.LpMaximize)

    x = {(k, i): pulp.LpVariable(f"x_{k}_{i}", cat="Binary") for k in ks for i in ids}
    s = {
        (k, i): pulp.LpVariable(f"s_{k}_{i}", cat="Binary")
        for k in ks
        for i in ids
        if pode_iniciar(k, i)
    }
    c = {
        (k, i, j): pulp.LpVariable(f"c_{k}_{i}_{j}", cat="Binary")
        for k in ks
        for (i, j) in conex
    }
    z = {i: pulp.LpVariable(f"z_{i}", cat="Binary") for i in ids}

    mdl += pulp.lpSum(miss[i]["prio"] * z[i] for i in ids), "missoes_cumpridas_ponderadas"

    for i in ids:
        mdl += z[i] == pulp.lpSum(x[(k, i)] for k in ks), f"def_z_{i}"

    for k in ks:
        for i in ids:
            entradas = pulp.lpSum(c[(k, jj, i)] for (jj, ii) in conex if ii == i)
            inicio = s[(k, i)] if (k, i) in s else 0
            mdl += inicio + entradas == x[(k, i)], f"in_{k}_{i}"

    for k in ks:
        for i in ids:
            saidas = pulp.lpSum(c[(k, i, jj)] for (ii, jj) in conex if ii == i)
            mdl += saidas <= x[(k, i)], f"out_{k}_{i}"

    for k in ks:
        inicios = [s[(k, i)] for i in ids if (k, i) in s]
        if inicios:
            mdl += pulp.lpSum(inicios) <= 1, f"start_{k}"

    for k in ks:
        mdl += (
            pulp.lpSum(miss[i]["dur"] * x[(k, i)] for i in ids)
            <= LIMITE_HORAS_F - acft[k]["f"]
        ), f"horas_{k}"

    solver = pulp.PULP_CBC_CMD(msg=False)
    mdl.solve(solver)
    prio_otima = int(round(pulp.value(mdl.objective) or 0))

    # Criterio secundario: entre solucoes com a mesma prioridade, minimiza horas
    # totais atribuidas. Neste v2 isso evita voos extras irrelevantes se houver
    # empate em prioridade.
    mdl += pulp.lpSum(miss[i]["prio"] * z[i] for i in ids) == prio_otima, "fixa_prioridade"
    mdl.sense = pulp.LpMinimize
    mdl.setObjective(pulp.lpSum(miss[i]["dur"] * x[(k, i)] for k in ks for i in ids))
    mdl.solve(solver)

    return {
        "modelo": mdl,
        "miss": miss,
        "acft": acft,
        "ids": ids,
        "ks": ks,
        "conex": conex,
        "x": x,
        "s": s,
        "c": c,
        "z": z,
        "horizonte_inicio": horizonte_inicio,
        "prioridade_otima": prio_otima,
    }


def rota_da_aeronave(sol: dict[str, object], k: str) -> list[str]:
    ids = sol["ids"]
    conex = sol["conex"]
    s = sol["s"]
    c = sol["c"]
    inicio = [
        i
        for i in ids
        if (k, i) in s and s[(k, i)].value() is not None and s[(k, i)].value() > 0.5
    ]
    if not inicio:
        return []

    seq = [inicio[0]]
    while True:
        atual = seq[-1]
        prox = [
            j
            for (ii, j) in conex
            if ii == atual and c[(k, atual, j)].value() is not None and c[(k, atual, j)].value() > 0.5
        ]
        if not prox:
            break
        seq.append(prox[0])
    return seq


def escrever_resultado(sol: dict[str, object], caminho: Path) -> None:
    caminho.parent.mkdir(parents=True, exist_ok=True)
    miss = sol["miss"]
    z = sol["z"]
    linhas = []
    for k in sol["ks"]:
        for ordem, mid in enumerate(rota_da_aeronave(sol, k), start=1):
            obj = miss[mid]["obj"]
            linhas.append(
                {
                    "aeronave_modelo": k,
                    "ordem": ordem,
                    "id": mid,
                    "om": obj.om,
                    "orig": obj.orig,
                    "dest": obj.dest,
                    "partida_utc": obj.partida_utc.isoformat().replace("+00:00", "Z"),
                    "chegada_utc": obj.chegada_utc.isoformat().replace("+00:00", "Z"),
                    "dur_h": f"{obj.dur_h:.2f}",
                    "prioridade": obj.prioridade,
                    "periodo": obj.periodo,
                    "aeronave_real_2025": obj.aeronave_real,
                    "linhas_origem": obj.linhas_origem,
                    "status": "cumprida",
                }
            )

    for mid in sol["ids"]:
        if not z[mid].value() or z[mid].value() <= 0.5:
            obj = miss[mid]["obj"]
            linhas.append(
                {
                    "aeronave_modelo": "",
                    "ordem": "",
                    "id": mid,
                    "om": obj.om,
                    "orig": obj.orig,
                    "dest": obj.dest,
                    "partida_utc": obj.partida_utc.isoformat().replace("+00:00", "Z"),
                    "chegada_utc": obj.chegada_utc.isoformat().replace("+00:00", "Z"),
                    "dur_h": f"{obj.dur_h:.2f}",
                    "prioridade": obj.prioridade,
                    "periodo": obj.periodo,
                    "aeronave_real_2025": obj.aeronave_real,
                    "linhas_origem": obj.linhas_origem,
                    "status": "nao_cumprida",
                }
            )

    colunas = [
        "status",
        "aeronave_modelo",
        "ordem",
        "id",
        "om",
        "orig",
        "dest",
        "partida_utc",
        "chegada_utc",
        "dur_h",
        "prioridade",
        "periodo",
        "aeronave_real_2025",
        "linhas_origem",
    ]
    with caminho.open("w", newline="", encoding="utf-8-sig") as arquivo:
        escritor = csv.DictWriter(arquivo, fieldnames=colunas)
        escritor.writeheader()
        escritor.writerows(linhas)


def imprimir_relatorio(sol: dict[str, object], resultado_csv: Path) -> None:
    miss = sol["miss"]
    z = sol["z"]
    ids = sol["ids"]
    cobertas = [i for i in ids if z[i].value() is not None and z[i].value() > 0.5]
    prioridade_total = sum(int(miss[i]["prio"]) for i in ids)
    prioridade_obtida = sum(int(miss[i]["prio"]) for i in cobertas)

    print("=" * 72)
    print(" OAMRP V2 - Registro de voo 2025 consolidado")
    print("=" * 72)
    print(f" Status solver: {pulp.LpStatus[sol['modelo'].status]}")
    print(f" Horizonte inicia em UTC: {sol['horizonte_inicio'].isoformat().replace('+00:00', 'Z')}")
    print(f" Missoes cumpridas: {len(cobertas)}/{len(ids)}")
    print(f" Prioridade obtida: {prioridade_obtida}/{prioridade_total}")
    print(f" Resultado CSV: {resultado_csv}")

    for k in sol["ks"]:
        seq = rota_da_aeronave(sol, k)
        horas = sum(float(miss[i]["dur"]) for i in seq)
        print(f"\n Aeronave {k}: {len(seq)} missoes | {horas:.1f}h")
        for mid in seq[:12]:
            obj = miss[mid]["obj"]
            print(
                f"   {obj.partida_utc:%d/%m %H:%M}Z "
                f"{obj.orig}->{obj.dest} OM {obj.om} ({mid})"
            )
        if len(seq) > 12:
            print(f"   ... mais {len(seq) - 12} missoes")

    nao_cumpridas = [i for i in ids if i not in cobertas]
    if nao_cumpridas:
        print("\n Missoes nao cumpridas (primeiras 20):")
        for mid in nao_cumpridas[:20]:
            obj = miss[mid]["obj"]
            print(f"   {obj.partida_utc:%d/%m %H:%M}Z {obj.orig}->{obj.dest} OM {obj.om} ({mid})")
        if len(nao_cumpridas) > 20:
            print(f"   ... mais {len(nao_cumpridas) - 20} missoes")


def main() -> None:
    parser = argparse.ArgumentParser(description="Resolve OAMRP v2 com registro de voo real.")
    parser.add_argument("--csv", type=Path, default=CSV_REGISTRO_PADRAO)
    parser.add_argument("--inicio", default=DATA_INICIO_PADRAO)
    parser.add_argument("--fim", default=DATA_FIM_PADRAO)
    parser.add_argument("--resultado", type=Path, default=CSV_RESULTADO_PADRAO)
    args = parser.parse_args()

    inicio = parse_data_inicio(args.inicio)
    fim = parse_data_inicio(args.fim)
    missoes = carregar_missoes(args.csv, inicio, fim)
    sol = resolver(missoes)
    escrever_resultado(sol, args.resultado)
    imprimir_relatorio(sol, args.resultado)


if __name__ == "__main__":
    main()