factorio-ai-agent/state_reader.py

"""
state_reader.py
RCON을 통해 팩토리오 게임 상태를 읽어오는 모듈

핵심 변경:
- game.player → game.players[1] (RCON 호환)
- scan_resources 반경 200→500 (초반 탐색 개선)
- get_buildings: name 대신 type 기반 검색 (Factorio 2.0 호환)
"""
import json
from factorio_rcon import FactorioRCON


P = """local p = game.players[1]
if not p then rcon.print("{}") return end
"""


class StateReader:
    def __init__(self, rcon: FactorioRCON):
        self.rcon = rcon

    def get_full_state(self) -> dict:
        return {
            "player":    self.get_player_info(),
            "inventory": self.get_inventory(),
            "resources": self.scan_resources(),
            "buildings": self.get_buildings(),
            "tech":      self.get_research_status(),
        }

    def get_player_info(self) -> dict:
        lua = P + """
rcon.print(game.table_to_json({
  x = math.floor(p.position.x),
  y = math.floor(p.position.y),
  health = p.character and p.character.health or 100
}))
"""
        raw = self.rcon.lua(lua)
        return json.loads(raw) if raw else {}

    def get_inventory(self) -> dict:
        lua = P + """
local inv = p.get_main_inventory()
if not inv then rcon.print("{}") return end
local result = {}
local contents = inv.get_contents()
for name, count in pairs(contents) do
  result[name] = count
end
rcon.print(game.table_to_json(result))
"""
        raw = self.rcon.lua(lua)
        return json.loads(raw) if raw else {}

    def scan_resources(self, radius: int = 500) -> dict:
        """플레이어 주변 자원 패치 위치 탐색 (반경 500타일)"""
        lua = P + f"""
local surface = p.surface
local center = p.position
local resources = {{}}

-- type="resource"로 모든 자원을 한번에 검색 (이름 호환 문제 없음)
local all_res = surface.find_entities_filtered{{
    area = {{
      {{center.x - {radius}, center.y - {radius}}},
      {{center.x + {radius}, center.y + {radius}}}
    }},
    type = "resource"
}}

-- 자원 이름별로 그룹핑
for _, e in ipairs(all_res) do
    local name = e.name
    if not resources[name] then
        resources[name] = {{count = 0, sx = 0, sy = 0}}
    end
    local r = resources[name]
    r.count = r.count + 1
    r.sx = r.sx + e.position.x
    r.sy = r.sy + e.position.y
end

-- 중심점 계산
local result = {{}}
for name, r in pairs(resources) do
    result[name] = {{
        count = r.count,
        center_x = math.floor(r.sx / r.count),
        center_y = math.floor(r.sy / r.count)
    }}
end
rcon.print(game.table_to_json(result))
"""
        raw = self.rcon.lua(lua)
        return json.loads(raw) if raw else {}

    def get_buildings(self, radius: int = 300) -> dict:
        """type 기반 검색으로 Factorio 2.0 호환"""
        lua = P + f"""
local surface = p.surface
local center = p.position
local area = {{
    {{center.x - {radius}, center.y - {radius}}},
    {{center.x + {radius}, center.y + {radius}}}
}}

-- type 기반으로 검색 (엔티티 이름 변경에 안전)
local types = {{
    "mining-drill", "furnace", "assembling-machine",
    "transport-belt", "inserter", "electric-pole",
    "generator", "boiler", "offshore-pump",
    "lab", "radar", "container",
    "ammo-turret", "electric-turret", "wall", "gate",
    "oil-refinery", "chemical-plant", "mining-drill"
}}

local result = {{}}
for _, t in ipairs(types) do
    local ok, entities = pcall(function()
        return surface.find_entities_filtered{{area = area, type = t}}
    end)
    if ok and entities and #entities > 0 then
        -- 이름별로 세분화
        for _, e in ipairs(entities) do
            result[e.name] = (result[e.name] or 0) + 1
        end
    end
end
rcon.print(game.table_to_json(result))
"""
        raw = self.rcon.lua(lua)
        return json.loads(raw) if raw else {}

    def get_research_status(self) -> dict:
        lua = P + """
local force = p.force
local completed = {}
local k = 0
for name, tech in pairs(force.technologies) do
  if tech.researched then
    k = k + 1
    completed[k] = name
  end
end
rcon.print(game.table_to_json({
  current = force.current_research and force.current_research.name or "none",
  completed_count = k
}))
"""
        raw = self.rcon.lua(lua)
        return json.loads(raw) if raw else {}

    def summarize_for_ai(self, state: dict) -> str:
        p   = state.get("player", {})
        inv = state.get("inventory", {})
        res = state.get("resources", {})
        bld = state.get("buildings", {})

        lines = [
            "## 현재 게임 상태",
            "### 플레이어",
            f"- 위치: ({p.get('x', '?')}, {p.get('y', '?')})",
            f"- 체력: {p.get('health', '?')}",
            "",
            "### 인벤토리 (주요 아이템)",
        ]

        key_items = [
            "iron-ore", "copper-ore", "coal", "stone",
            "iron-plate", "copper-plate", "steel-plate",
            "burner-mining-drill", "electric-mining-drill",
            "stone-furnace", "transport-belt", "inserter",
            "burner-inserter", "small-electric-pole",
        ]
        for item in key_items:
            count = inv.get(item, 0)
            if count > 0:
                lines.append(f"- {item}: {count}개")

        if not any(inv.get(item, 0) > 0 for item in key_items):
            lines.append("- 비어 있음")

        lines += ["", "### 주변 자원 패치 (반경 500타일 스캔)"]
        if res:
            # 거리순으로 정렬
            px = p.get('x', 0)
            py = p.get('y', 0)
            sorted_res = sorted(
                res.items(),
                key=lambda item: ((item[1]['center_x'] - px)**2 + (item[1]['center_y'] - py)**2)
            )
            for ore, info in sorted_res:
                dist = int(((info['center_x'] - px)**2 + (info['center_y'] - py)**2)**0.5)
                lines.append(
                    f"- {ore}: {info['count']}타일 "
                    f"(중심: {info['center_x']}, {info['center_y']}) "
                    f"[거리: ~{dist}타일]"
                )
        else:
            lines.append("- 반경 500타일 내 자원 없음 — 더 멀리 탐색 필요")

        lines += ["", "### 건설된 건물"]
        if bld:
            for name, count in bld.items():
                lines.append(f"- {name}: {count}개")
        else:
            lines.append("- 아직 건물 없음")

        return "\n".join(lines)