"""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.4"

import warnings

import httpx


def _check_version_compat(client_version: str, server_version: str) -> None:
    """Warn if client and server minor versions differ."""
    try:
        c_major, c_minor, _ = client_version.split(".")
        s_major, s_minor, _ = server_version.split(".")
    except (ValueError, AttributeError):
        return
    if c_major != s_major:
        warnings.warn(
            f"gravi-model-client v{client_version} may be incompatible with "
            f"server (expects client v{server_version}). "
            f"Major version mismatch — please upgrade.",
            stacklevel=3,
        )
    elif c_minor != s_minor:
        warnings.warn(
            f"gravi-model-client v{client_version} may be incompatible with "
            f"server (expects client v{server_version}). "
            f"Consider upgrading: pip install --upgrade gravi-model-client",
            stacklevel=3,
        )


class GraviModelClient:
    """Async 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.AsyncClient(
            headers={"X-API-Key": api_key},
            timeout=timeout,
        )

    async def __aenter__(self) -> "GraviModelClient":
        return self

    async def __aexit__(self, *args: object) -> None:
        await self.close()

    async def close(self) -> None:
        await self.client.aclose()

    async def health_check(self) -> dict:
        response = await self.client.get(f"{self.base_url}/healthz")
        response.raise_for_status()
        data = response.json()
        if server_version := data.get("client_version"):
            _check_version_compat(__version__, server_version)
        return data

    async def run_demand_model(self, request: ModelRunRequest) -> ModelResultResponse:
        response = await 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())

    async def run_tsd_model(self, optimization: OptimizationInput) -> OptimizationAPIResponse:
        response = await 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",
]
