import asyncio
from contextlib import contextmanager
from typing import Dict
import httpx
import pandas as pd
from httpx import Response


class DTNClient(httpx.AsyncClient):
    def __init__(
        self,
        username: str,
        web_service_key: str,
        api_key: str,
        base_url: str,
        timeout: float = 180.0,
    ):
        super().__init__(base_url=base_url)
        self.username = username
        self.web_service_key = web_service_key
        self.api_key = api_key
        self.base_url = base_url.rstrip("/")
        self._base_url_str = base_url.rstrip("/")
        self.client = httpx.AsyncClient(timeout=180)
        self.timeout = httpx.Timeout(timeout, connect=60.0)

    def __repr__(self):
        return """DTN API Client designed for DTN Allocation Tracker"""

    @property
    def custom_headers(self) -> Dict[str, str]:
        return {
            "username": self.username,
            "Accept": "application/vnd.dtn.energy.v1+JSON",
            "webservicekey": self.web_service_key,
            "apikey": self.api_key,
        }

    async def get_allocations(self) -> Response:
        url = "/allocations"
        res = await self.get(url=url, headers=self.custom_headers, timeout=self.timeout)
        res.raise_for_status()
        return res

    async def get_data_in_group(self, url: str) -> Response:
        if url.startswith(self._base_url_str):
            url = url.removeprefix(self._base_url_str)
        if not url.startswith("/"):
            url = f"/{url}"
        res = await self.get(url=url, headers=self.custom_headers, timeout=self.timeout)
        return res


@contextmanager
def init_dtn_api(username=None, web_service_key=None, api_key=None) -> DTNClient:
    dtn_client = DTNClient(
        username=username,
        web_service_key=web_service_key,
        api_key=api_key,
        base_url="https://api.dtn.com/fuelsuite/allocationtracker/",
    )
    yield dtn_client


def parse_allocation_data(allocation_data):
    ret = []
    for allocation in allocation_data["data"]:
        supplier = allocation.get("supplier", {})
        location_data = allocation.get("location", {})
        customer_data = allocation.get("customer", {})
        consignee_group_data = customer_data.get("consigneeGroup", {})
        product_group_data = allocation.get("productAllocationList", [])
        terminal_data = location_data.get("terminal", {}) or {}
        location_group_data = location_data.get("terminalGroup", {}) or {}
        sold_to_data = consignee_group_data.get("soldTo", {}) or {}
        supplier_name = supplier.get("name")
        supplier_code = supplier.get("sellerNum")
        terminal_name = (
            location_group_data.get("name") or terminal_data.get("name") or None
        )
        terminal_code = location_group_data.get("id") or terminal_data.get("id") or None
        terminal_hfref = location_group_data.get("terminalListLink", {}).get("href")
        consignee_name = consignee_group_data.get("name")
        consignee_id = consignee_group_data.get("id")
        sold_to_id = sold_to_data.get("id")
        sold_to_name = sold_to_data.get("name")
        for group in product_group_data:
            product_group_name = group.get("name")
            uom = group.get("unitOfMeasure")
            allocation_remaining_monthly = group.get("allocationRemainingAmountMonthly")
            lifted_amount_monthly = group.get("liftedAmountMonthly")
            refresh_amount_monthly = group.get("refreshAmountMonthly")
            refresh_date_monthly = group.get("refreshDateMonthly")
            scaled_start_amount_monthly = group.get("scaledStartAmountMonthly")
            start_amount_monthly = group.get("startAmountMonthly")
            row = {
                "supplier_name": supplier_name,
                "supplier_code": supplier_code,
                "terminal_name": terminal_name,
                "terminal_code": terminal_code,
                "terminal_hfref": terminal_hfref,
                "consignee_name": consignee_name,
                "consignee_id": consignee_id,
                "sold_to_name": sold_to_name,
                "sold_to_id": sold_to_id,
                "product_group_name": product_group_name,
                "uom": uom,
                "allocation_remaining_monthly": allocation_remaining_monthly,
                "lifted_amount_monthly": lifted_amount_monthly,
                "refresh_amount_monthly": refresh_amount_monthly,
                "refresh_date_monthly": refresh_date_monthly,
                "scaled_start_amount_monthly": scaled_start_amount_monthly,
                "start_amount_monthly": start_amount_monthly,
            }
            ret.append(row)
    return pd.DataFrame(ret)


if __name__ == "__main__":

    async def get_allocation_data():
        with init_dtn_api(username="", web_service_key="", api_key="") as dtn_client:
            allocations = await dtn_client.get_allocations()
            print(allocations)

    asyncio.run(get_allocation_data())
