api/Api.py

import requests
from datetime import datetime
import url64
import streamlit as st
from config import CONFIG
from excel.Formatter import RoutePlan
from .queries import CREATE_COLLECTION, DISABLE_COLLECTIONS_BY_DATE

class EstimatedContainer:
    def __init__(self, amount: int, material_type_id: int, container_type_id: int, weight_kg: int | None):
        self.amount = amount
        self.material_type_id = material_type_id
        self.container_type_id = container_type_id
        self.weight_kg = weight_kg

    def to_dict(self):
        return {
            "amount": self.amount,
            "materialTypeId": self.material_type_id,
            "containerTypeId": self.container_type_id,
            "weightKg": self.weight_kg
        }
    
class Collection:
    def __init__(self, branch_id: int, driver_id: int, vehicle_id: int, estimated_containers: list[EstimatedContainer]):
        self.driver_id            = driver_id
        self.vehicle_id           = vehicle_id
        self.branch_id            = branch_id
        self.estimated_containers = estimated_containers

class Route:
    def __init__(self, id: int, ordered_collections: list[Collection]):
        self.id = id
        self.ordered_collections = ordered_collections

class Api:
    _API_V1_BASE_URL = CONFIG.API_V1_BASE_URL
    _API_V2_BASE_URL = CONFIG.API_V2_BASE_URL
    _API_AUTH_URL    = CONFIG.AUTH_URL
    _DEFAULT_HEADERS = {'Content-Type': 'application/json; charset=utf-8' }
    
    _DISABLE_COLLECTIONS_BY_DATE_QUERY = DISABLE_COLLECTIONS_BY_DATE
    _CREATE_COLLECTION_QUERY           = CREATE_COLLECTION

    @staticmethod
    def _get_auth_token(query_param_name="t") -> str:
        # Obtener el token en crudo (en formato base64url) del query param 't'
        query_params = st.experimental_get_query_params()
        raw_token = query_params.get(query_param_name, [None])[0]

        # Validar the el token esté presente
        if raw_token is None or raw_token == "":
            raise Exception("Query param t not found")
        
        # Quitar formato base64url
        token = url64.decode(raw_token)
        return token
    
    @staticmethod
    def _format_headers(token: str) -> dict:
        headers = dict(Api._DEFAULT_HEADERS)
        headers["Authorization"] = f"Bearer {token}"
        return headers

    @staticmethod
    def _get_recycler_id(token: str) -> int: 
        res = requests.post(Api._API_AUTH_URL, json={"authToken": token})
        print(f"[AUTH] Status: {res.status_code}")
        
        recycler_id = res.json().get("idClienteReciclador", None)
        if recycler_id is None:
            raise Exception("AUTHENTICATION FAILED")
        
        return int(recycler_id)

    def __init__(self):
        self.token       = Api._get_auth_token()
        self.recycler_id = Api._get_recycler_id(self.token)
        self.headers     = Api._format_headers(self.token)

    def _create_collection(self, route_id: int, driver_id: int, vehicle_id: int, branch_id: int, date: datetime, estimated_containers: list[EstimatedContainer]) -> int:
        res =  requests.post(Api._API_V2_BASE_URL, headers=self.headers, json={
            "query": Api._CREATE_COLLECTION_QUERY,
            "variables": {
                "route_id"            : int(route_id),
                "recycler_id"         : int(self.recycler_id),
                "driver_id"           : int(driver_id),
                "vehicle_id"          : int(vehicle_id),
                "branch_id"           : int(branch_id),
                "date"                : date.strftime("%Y-%m-%d"),
                "estimated_containers": [ec.to_dict() for ec in estimated_containers]
            }
        })
        print(f"[CREATE_COLLECTION] Status: {res.status_code}")
        collection_id = res.json()["data"]["collection"]["id"]
        if res.status_code != 200 or collection_id is None:
            print("[ERROR] Input: ", {
                "route_id": route_id, "driver_id": driver_id, "vehicle_id": vehicle_id, "branch_id": branch_id, "date": datetime, "estimated_containers": estimated_containers
            })
            print("[ERROR] Output: ", res.json())
            raise Exception("Failed to create collection")
        
        return collection_id

    # def _create_collection(self, branch_id: int, driver_id: int, date: datetime, estimated_containers: list[EstimatedContainer]) -> int:
    #     payload = {
    #         "id_sucursal"        : branch_id,
    #         "id_driver"          : driver_id,
    #         "date"               : str(date),
    #         "schedule_hr"        : [],
    #         "materials"          : [],
    #         "estimate_weight_kg" : [],
    #         "estimate_volume_m3" : None,
    #         "estimate_containers": [ec.to_dict() for ec in estimated_containers],
    #         "priority"           : 0
    #     }


    #     # Crear nueva recolección
    #     res = requests.post(f"{Api._API_V1_BASE_URL}/clientes-recicladores/{self.recycler_id}/recollections", headers=self.headers, json=payload)
    #     print(f"[CREATE_COLLECTION] Status: {res.status_code}")
        
    #     # Validar resultado
    #     collection_id = res.json().get("id", None)
    #     if res.status != 201 or res.status != 200 or collection_id is None:
    #         raise Exception("Failed to create recollection with input:", payload)

    #     return collection_id
    
    def _reorder_route(self, route_id: int, ordered_collection_ids: list[int]):
        """
        collections is in format (collection_id, index)
        """

        payload: list[dict] = []
        for i in range(0, len(ordered_collection_ids)):
            id = ordered_collection_ids[i]
            payload.append({"id_recollection": id, "index": i})

        res = requests.put(f"{Api._API_V1_BASE_URL}/clientes-recicladores/{self.recycler_id}/routes-order/{route_id}", headers=self.headers, json=payload)
        print(f"[REORDER_ROUTES] Status: {res.status_code}")
        
        if res.status_code != 200:
            raise Exception("Failed to update route order")

    def _disable_collections_by_date(self, date: datetime, excluded_collection_ids: list[int]):
        res = requests.post(Api._API_V2_BASE_URL, headers=self.headers, json={
            "query": Api._DISABLE_COLLECTIONS_BY_DATE_QUERY,
            "variables": {
                "recycler_id" : self.recycler_id,
                "date"        : str(date),
                "excluded_ids": excluded_collection_ids
            }
        })
        print(f"[DISABLE_COLLECTIONS] Status: {res.status_code}")
        
        if res.status_code != 200:
            raise Exception("Failed to disable collections by date")

    def modify_daily_routes(self, date: datetime, routes: list[Route]) -> list[int]:
        """
        Returns the ids of the newly created collections.
        @TODO: Make this transaction-like.
        """


        # Crear recolecciones nuevas del día indicado
        created_collection_ids: list[int] = []
        for r in routes:
            ordered_collection_ids: list[int] = [] 
            for c in r.ordered_collections:
                new_collection_id = self._create_collection(
                    route_id=r.id,
                    vehicle_id=c.vehicle_id,
                    branch_id=c.branch_id,
                    driver_id=c.driver_id,
                    date=date,
                    estimated_containers=c.estimated_containers
                )
                ordered_collection_ids.append(new_collection_id)
                created_collection_ids.append(new_collection_id)
            
            # Ordenar recolecciones después de crearlas
            self._reorder_route(route_id=r.id, ordered_collection_ids=ordered_collection_ids)
        
        # Deshabilitar todas las recolecciones del mismo día
        self._disable_collections_by_date(date=date, excluded_collection_ids=created_collection_ids)

        return created_collection_ids