from pydantic import BaseModel, Field


class VastPlatformQuery(BaseModel):
    bw_nvlink: float | None = Field(default=None, description="Bandwidth NVLink")
    compute_cap: int | None = Field(default=None, description="CUDA compute capability * 100 (ie: 650 for 6.5, 700 for 7.0)")
    cpu_cores: int | None = Field(default=None, description="Number of virtual CPU cores")
    cpu_cores_effective: float | None = Field(default=None, description="Effective number of CPU cores")
    cpu_ram: float | None = Field(default=None, description="System RAM in gigabytes")
    cuda_vers: float | None = Field(default=None, description="Machine max supported CUDA version (based on driver version)")
    datacenter: bool | None = Field(default=None, description="Show only datacenter offers")
    direct_port_count: int | None = Field(default=None, description="Open ports on host's router")
    disk_bw: float | None = Field(default=None, description="Disk read bandwidth, in MB/s")
    disk_space: float | None = Field(default=None, description="Disk storage space, in GB")
    dlperf: float | None = Field(default=None, description="DL-perf score")
    dlperf_usd: float | None = Field(default=None, description="DL-perf/$")
    dph: float | None = Field(default=None, description="$/hour rental cost")
    driver_version: str | None = Field(default=None, description="Machine's NVIDIA driver version as 3 digit string ex. '535.86.05'")
    duration: float | None = Field(default=None, description="Max rental duration in days")
    external: bool | None = Field(default=False, description="Show external offers in addition to datacenter offers")
    flops_usd: float | None = Field(default=None, description="TFLOPs/$")
    geolocation: str | None = Field(default=None, description="Two letter country code. Works with operators =, !=, in, not in (e.g. geolocation not in [XV,XZ])")
    gpu_mem_bw: float | None = Field(default=None, description="GPU memory bandwidth in GB/s")
    gpu_name: str | None = Field(default=None, description="GPU model name (no quotes, replace spaces with underscores, ie: RTX_3090 rather than 'RTX 3090')")
    gpu_ram: float | None = Field(default=None, description="GPU RAM in GB")
    gpu_frac: float | None = Field(default=None, description="Ratio of GPUs in the offer to GPUs in the system")
    gpu_display_active: bool | None = Field(default=None, description="True if the GPU has a display attached")
    has_avx: bool | None = Field(default=None, description="CPU supports AVX instruction set")
    id: int | None = Field(default=None, description="Instance unique ID")
    inet_down: float | None = Field(default=None, description="Internet download speed in Mb/s")
    inet_down_cost: float | None = Field(default=None, description="Internet download bandwidth cost in $/GB")
    inet_up: float | None = Field(default=None, description="Internet upload speed in Mb/s")
    inet_up_cost: float | None = Field(default=None, description="Internet upload bandwidth cost in $/GB")
    machine_id: int | None = Field(default=None, description="Machine ID of instance")
    min_bid: float | None = Field(default=None, description="Current minimum bid price in $/hr for interruptible")
    num_gpus: int | None = Field(default=None, description="# of GPUs")
    pci_gen: float | None = Field(default=None, description="PCIE generation")
    pcie_bw: float | None = Field(default=None, description="PCIE bandwidth")
    reliability: float | None = Field(default=None, description="Machine reliability score")
    rentable: bool | None = Field(default=True, description="Is the instance currently rentable")
    rented: bool | None = Field(default=False, description="Allow/disallow duplicates and potential conflicts with existing stopped instances")
    storage_cost: float | None = Field(default=None, description="Storage cost in $/GB/month")
    static_ip: bool | None = Field(default=None, description="Is the IP address static/stable")
    total_flops: float | None = Field(default=None, description="Total TFLOPs from all GPUs")
    verified: bool | None = Field(default=True, description="Is the machine verified")


class VastPlatformSettings(BaseModel):
    instance_id: int | None = None
    advanced_query_settings: VastPlatformQuery | None = Field(default=None, description="Configure search settings to find GPU instances on-demand")
    disk_size: int | None = Field(default=None, description="Disk storage space, in GB")
