"""Gravi-Model Python client for demand and supply optimization."""

from gravi_model_client._base import CamelModel
from gravi_model_client.demand import (
    CostDetails,
    DropResponse,
    ERVolume,
    InitialVolume,
    InventoryRecord,
    ModelInput,
    ModelResultResponse,
    ModelRunRequest,
    Order,
    OrderResponse,
    StoreRequest,
    TankInventory,
    TankPeriod,
    TankRequest,
    TankVolume,
    UnscheduledVolume,
    VolumeRange,
)
from gravi_model_client.tsd import (
    Blend,
    BlendComponent,
    BlendComponentDetail,
    ComponentDetail,
    Config,
    ContractAssignment,
    ContractUtilization,
    Directive,
    DirectiveSupplyPrice,
    FakeSupplyUsage,
    OptimizationAPIResponse,
    OptimizationInput,
    OptimizationMetadata,
    OptimizationResult,
    PointToPoint,
    PrimaryTankMode,
    Product,
    RackPrice,
    Site,
    SiteAssignment,
    SupplyPrice,
    Tank,
    TankAssignmentDetail,
    TerminalLink,
)

__version__ = "0.1.1"

import httpx


class GraviModelClient:
    """Client for Gravitate demand and supply optimization API."""

    def __init__(
        self,
        api_key: str,
        base_url: str = "https://demand-model.gravitate.energy",
        timeout: float = 120.0,
    ) -> None:
        self.api_key = api_key
        self.base_url = base_url.rstrip("/")
        self.client = httpx.Client(
            headers={"X-API-Key": api_key},
            timeout=timeout,
        )

    def __enter__(self) -> "GraviModelClient":
        return self

    def __exit__(self, *args: object) -> None:
        self.close()

    def close(self) -> None:
        self.client.close()

    def health_check(self) -> dict:
        response = self.client.get(f"{self.base_url}/health")
        response.raise_for_status()
        return response.json()

    def run_demand_model(self, request: ModelRunRequest) -> ModelResultResponse:
        response = self.client.post(
            f"{self.base_url}/run",
            content=request.model_dump_json(by_alias=True),
            headers={"Content-Type": "application/json"},
        )
        response.raise_for_status()
        return ModelResultResponse.model_validate(response.json())

    def run_tsd_model(self, optimization: OptimizationInput) -> OptimizationAPIResponse:
        response = self.client.post(
            f"{self.base_url}/tsd/optimize",
            content=optimization.model_dump_json(),
            headers={"Content-Type": "application/json"},
        )
        response.raise_for_status()
        return OptimizationAPIResponse.model_validate(response.json())


__all__ = [
    "CamelModel",
    "GraviModelClient",
    # Demand models
    "CostDetails",
    "DropResponse",
    "ERVolume",
    "InitialVolume",
    "InventoryRecord",
    "ModelInput",
    "ModelResultResponse",
    "ModelRunRequest",
    "Order",
    "OrderResponse",
    "StoreRequest",
    "TankInventory",
    "TankPeriod",
    "TankRequest",
    "TankVolume",
    "UnscheduledVolume",
    "VolumeRange",
    # TSD models
    "Blend",
    "BlendComponent",
    "BlendComponentDetail",
    "ComponentDetail",
    "Config",
    "ContractAssignment",
    "ContractUtilization",
    "Directive",
    "DirectiveSupplyPrice",
    "FakeSupplyUsage",
    "OptimizationAPIResponse",
    "OptimizationInput",
    "OptimizationMetadata",
    "OptimizationResult",
    "PointToPoint",
    "PrimaryTankMode",
    "Product",
    "RackPrice",
    "Site",
    "SiteAssignment",
    "SupplyPrice",
    "Tank",
    "TankAssignmentDetail",
    "TerminalLink",
]
