backend/game/units.py

from __future__ import annotations

import uuid
from enum import Enum
from typing import Optional

from pydantic import BaseModel, Field


class UnitType(str, Enum):
    SCV = "scv"
    MARINE = "marine"
    MEDIC = "medic"
    GOLIATH = "goliath"
    TANK = "tank"
    WRAITH = "wraith"


class UnitStatus(str, Enum):
    IDLE = "idle"
    MOVING = "moving"
    ATTACKING = "attacking"
    MINING_MINERALS = "mining_minerals"
    MINING_GAS = "mining_gas"
    MOVING_TO_BUILD = "moving_to_build"
    BUILDING = "building"
    HEALING = "healing"
    SIEGED = "sieged"
    PATROLLING = "patrolling"


class UnitDef(BaseModel):
    max_hp: int
    armor: int
    ground_damage: int
    air_damage: int
    attack_range: float
    move_speed: float          # tiles/second
    mineral_cost: int
    gas_cost: int
    supply_cost: int
    build_time_ticks: int      # ticks to produce (at 4 ticks/s)
    is_flying: bool = False
    can_cloak: bool = False
    can_siege: bool = False
    heal_per_tick: float = 0.0 # medic healing
    attack_cooldown_ticks: int = 4
    # Siege-mode stats (tanks only)
    siege_damage: int = 0
    siege_range: float = 0.0
    siege_splash_radius: float = 0.0
    siege_cooldown_ticks: int = 0


UNIT_DEFS: dict[UnitType, UnitDef] = {
    UnitType.SCV: UnitDef(
        max_hp=60, armor=0, ground_damage=5, air_damage=0,
        attack_range=1, move_speed=1.5, mineral_cost=50, gas_cost=0,
        supply_cost=1, build_time_ticks=20, attack_cooldown_ticks=6,
    ),
    UnitType.MARINE: UnitDef(
        max_hp=40, armor=0, ground_damage=6, air_damage=6,
        attack_range=4, move_speed=1.5, mineral_cost=50, gas_cost=0,
        supply_cost=1, build_time_ticks=24, attack_cooldown_ticks=4,
    ),
    UnitType.MEDIC: UnitDef(
        max_hp=60, armor=1, ground_damage=0, air_damage=0,
        attack_range=2, move_speed=1.5, mineral_cost=50, gas_cost=25,
        supply_cost=1, build_time_ticks=24, heal_per_tick=1.5,
        attack_cooldown_ticks=999,  # medics don't attack
    ),
    UnitType.GOLIATH: UnitDef(
        max_hp=125, armor=1, ground_damage=12, air_damage=20,
        attack_range=5, move_speed=1.0, mineral_cost=100, gas_cost=50,
        supply_cost=2, build_time_ticks=40, attack_cooldown_ticks=4,
    ),
    UnitType.TANK: UnitDef(
        max_hp=150, armor=1, ground_damage=35, air_damage=0,
        attack_range=7, move_speed=0.75, mineral_cost=150, gas_cost=100,
        supply_cost=3, build_time_ticks=50,
        attack_cooldown_ticks=5,
    ),
    UnitType.WRAITH: UnitDef(
        max_hp=120, armor=0, ground_damage=8, air_damage=20,
        attack_range=5, move_speed=2.5, mineral_cost=150, gas_cost=100,
        supply_cost=2, build_time_ticks=60, is_flying=True,
        attack_cooldown_ticks=4,
    ),
}


class Unit(BaseModel):
    id: str = Field(default_factory=lambda: str(uuid.uuid4())[:8])
    unit_type: UnitType
    owner: str
    x: float
    y: float
    hp: float
    max_hp: int
    status: UnitStatus = UnitStatus.IDLE
    # Movement
    target_x: Optional[float] = None
    target_y: Optional[float] = None
    path_waypoints: list[list[float]] = Field(default_factory=list)  # [[x,y], ...] for pathfinding
    # Patrol waypoints (current + return)
    patrol_x: Optional[float] = None
    patrol_y: Optional[float] = None
    # Combat
    attack_target_id: Optional[str] = None
    attack_target_building_id: Optional[str] = None
    attack_cooldown: int = 0
    # Modes
    is_sieged: bool = False
    is_cloaked: bool = False
    # SCV tasks
    assigned_resource_id: Optional[str] = None
    building_target_id: Optional[str] = None
    harvest_carry: bool = False   # True while carrying resources back to CC
    harvest_amount: int = 0       # amount being carried (deposited on CC arrival)
    harvest_mining_ticks: int = 0  # ticks spent drilling at the current patch (0 = not yet mining)
    stuck_ticks: int = 0          # consecutive ticks the unit couldn't move at all
    nav_stall_ticks: int = 0      # ticks making no progress toward current waypoint (oscillation)

    @classmethod
    def create(cls, unit_type: UnitType, owner: str, x: float, y: float) -> "Unit":
        defn = UNIT_DEFS[unit_type]
        return cls(
            unit_type=unit_type,
            owner=owner,
            x=x,
            y=y,
            hp=float(defn.max_hp),
            max_hp=defn.max_hp,
        )

    def dist_to(self, x: float, y: float) -> float:
        return ((self.x - x) ** 2 + (self.y - y) ** 2) ** 0.5