factorio-ai-agent/context_compressor.py

"""
context_compressor.py — 계층적 컨텍스트 압축

핵심 변경:
- game.player → game.players[1] (RCON 호환)
- pcall로 엔티티 검색 감싸기 (Factorio 2.0 이름 변경 대응)
- type 기반 검색으로 버전 호환성 확보
"""
import json
import math
from factorio_rcon import FactorioRCON

ZONE_SIZE = 64

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


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

    def get_compressed_state(self, detail_level: int = 1) -> str:
        global_summary = self._get_global_summary()
        problems       = self._detect_problems()
        player         = self._get_player_info()

        if detail_level == 0:
            return self._format_global(global_summary, problems, player)

        zones = self._get_zone_summaries()

        if detail_level == 1:
            return self._format_level1(global_summary, zones, problems, player)

        drilldown = self._drilldown_problem_zones(zones, problems)
        return self._format_level2(global_summary, zones, problems, drilldown, player)

    # ── 글로벌 지표 (pcall로 안전하게) ──────────────────────────────

    def _get_global_summary(self) -> dict:
        lua = P + """
local surface = p.surface
local force   = p.force
local center  = p.position

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

local counts = {}
local total = 0
for _, t in ipairs(type_list) do
  local ok, found = pcall(function()
    return surface.find_entities_filtered{
      area={{center.x-500,center.y-500},{center.x+500,center.y+500}},
      type=t
    }
  end)
  if ok and found then
    for _, e in ipairs(found) do
      counts[e.name] = (counts[e.name] or 0) + 1
      total = total + 1
    end
  end
end

-- 연구 진행
local current_tech = "none"
local tech_progress = 0
if force.current_research then
  current_tech = force.current_research.name
  tech_progress = math.floor(force.research_progress * 100)
end

local researched = 0
for _, t in pairs(force.technologies) do
  if t.researched then researched = researched + 1 end
end

rcon.print(game.table_to_json({
  total_entities = total,
  counts = counts,
  current_research = current_tech,
  research_progress = tech_progress,
  researched_count = researched,
  evolution = math.floor(force.evolution_factor * 100)
}))
"""
        raw = self.rcon.lua(lua)
        try:
            return json.loads(raw) if raw else {}
        except Exception:
            return {}

    # ── 구역 요약 ───────────────────────────────────────────────────

    def _get_zone_summaries(self) -> list[dict]:
        lua = P + f"""
local surface = p.surface
local center  = p.position
local Z = {ZONE_SIZE}
local radius = 512
local zones  = {{}}

local function zone_key(x, y)
  return math.floor(x/Z) .. "," .. math.floor(y/Z)
end

local all_entities = surface.find_entities_filtered{{
  area={{{{center.x-radius, center.y-radius}},
         {{center.x+radius, center.y+radius}}}},
  force = "player"
}}

for _, e in ipairs(all_entities) do
  local key = zone_key(e.position.x, e.position.y)
  if not zones[key] then
    zones[key] = {{
      key = key,
      x   = math.floor(e.position.x/Z)*Z,
      y   = math.floor(e.position.y/Z)*Z,
      entities = 0,
      miners = 0, furnaces = 0, assemblers = 0,
      belts = 0, inserters = 0, poles = 0,
      turrets = 0, labs = 0, pumps = 0
    }}
  end
  local z = zones[key]
  z.entities = z.entities + 1
  local n = e.name
  local t = e.type
  if t == "mining-drill" then z.miners = z.miners + 1
  elseif t == "furnace" then z.furnaces = z.furnaces + 1
  elseif t == "assembling-machine" then z.assemblers = z.assemblers + 1
  elseif t == "transport-belt" then z.belts = z.belts + 1
  elseif t == "inserter" then z.inserters = z.inserters + 1
  elseif t == "electric-pole" then z.poles = z.poles + 1
  elseif t == "ammo-turret" or t == "electric-turret" then z.turrets = z.turrets + 1
  elseif n == "lab" then z.labs = z.labs + 1
  elseif t == "offshore-pump" then z.pumps = z.pumps + 1
  end
end

local result = {{}}
for _, z in pairs(zones) do
  if z.entities > 2 then
    result[#result+1] = z
  end
end
rcon.print(game.table_to_json(result))
"""
        raw = self.rcon.lua(lua)
        try:
            return json.loads(raw) if raw else []
        except Exception:
            return []

    # ── 문제 감지 ───────────────────────────────────────────────────

    def _detect_problems(self) -> list[str]:
        lua = P + """
local surface = p.surface
local force   = p.force
local center  = p.position
local problems = {}

-- 1. 연료 없는 버너 건물
local burner_entities = surface.find_entities_filtered{
  area={{center.x-300,center.y-300},{center.x+300,center.y+300}},
  type={"mining-drill","furnace"}
}
local no_fuel = 0
for _, e in ipairs(burner_entities) do
  if e.burner then
    local fi = e.burner.inventory
    if fi and fi.is_empty() then
      no_fuel = no_fuel + 1
    end
  end
end
if no_fuel > 0 then
  problems[#problems+1] = "연료 부족: " .. no_fuel .. "개 건물이 연료 없음"
end

-- 2. 꽉 찬 삽입기 (병목)
local inserters = surface.find_entities_filtered{
  area={{center.x-300,center.y-300},{center.x+300,center.y+300}},
  type="inserter"
}
local stuck = 0
for _, ins in ipairs(inserters) do
  if ins.held_stack and ins.held_stack.valid_for_read then
    stuck = stuck + 1
  end
end
if stuck > 5 then
  problems[#problems+1] = "벨트 병목: 삽입기 " .. stuck .. "개가 아이템 들고 멈춤"
end

-- 3. 자원 고갈 임박
local drills = surface.find_entities_filtered{
  area={{center.x-300,center.y-300},{center.x+300,center.y+300}},
  type="mining-drill"
}
local depleting = 0
for _, d in ipairs(drills) do
  local ore = surface.find_entities_filtered{
    area={{d.position.x-3, d.position.y-3},
          {d.position.x+3, d.position.y+3}},
    type="resource"
  }
  if #ore > 0 and ore[1].amount < 5000 then
    depleting = depleting + 1
  end
end
if depleting > 0 then
  problems[#problems+1] = "자원 고갈 임박: " .. depleting .. "개 채굴기 구역 5000 미만"
end

rcon.print(game.table_to_json(problems))
"""
        raw = self.rcon.lua(lua)
        try:
            return json.loads(raw) if raw else []
        except Exception:
            return []

    def _drilldown_problem_zones(self, zones: list, problems: list) -> str:
        if not problems or not zones:
            return ""
        top = sorted(zones, key=lambda z: z.get("entities", 0), reverse=True)[:2]
        lines = ["#### 주요 구역 상세"]
        for z in top:
            lines.append(
                f"  구역({z['x']},{z['y']}): "
                f"채굴기{z.get('miners',0)} 제련소{z.get('furnaces',0)} "
                f"조립기{z.get('assemblers',0)} 벨트{z.get('belts',0)} "
                f"삽입기{z.get('inserters',0)} 전선주{z.get('poles',0)}"
            )
        return "\n".join(lines)

    def _get_player_info(self) -> dict:
        lua = P + """
local inv = p.get_main_inventory()
if not inv then rcon.print("{}") return end
local contents = inv.get_contents()

local inv_summary = {}
for name, count in pairs(contents) do
  if count > 0 then inv_summary[name] = count end
end

rcon.print(game.table_to_json({
  x = math.floor(p.position.x),
  y = math.floor(p.position.y),
  inventory = inv_summary
}))
"""
        raw = self.rcon.lua(lua)
        try:
            return json.loads(raw) if raw else {}
        except Exception:
            return {}

    # ── 포맷터 ──────────────────────────────────────────────────────

    def _format_global(self, g: dict, problems: list, player: dict) -> str:
        lines = [
            "## 공장 상태 (글로벌 요약)",
            f"위치: ({player.get('x','?')}, {player.get('y','?')})",
            f"전체 건물 수: {g.get('total_entities', 0)}개",
            f"진화도: {g.get('evolution', 0)}%",
            f"완료 연구: {g.get('researched_count', 0)}개",
            f"진행 중 연구: {g.get('current_research','없음')} ({g.get('research_progress',0)}%)",
        ]
        if problems:
            lines.append("\n## 감지된 문제")
            for p in problems:
                lines.append(f"  ! {p}")
        lines += self._format_inventory(player.get("inventory", {}))
        return "\n".join(lines)

    def _format_level1(self, g: dict, zones: list, problems: list, player: dict) -> str:
        lines = [self._format_global(g, problems, player), "\n## 구역별 요약"]
        for z in sorted(zones, key=lambda x: x.get("entities", 0), reverse=True)[:8]:
            role = self._guess_zone_role(z)
            lines.append(
                f"  [{z['x']:+d},{z['y']:+d}] {role} | "
                f"채굴{z.get('miners',0)} 제련{z.get('furnaces',0)} "
                f"조립{z.get('assemblers',0)} 벨트{z.get('belts',0)}"
            )
        return "\n".join(lines)

    def _format_level2(self, g, zones, problems, drilldown, player) -> str:
        base = self._format_level1(g, zones, problems, player)
        if drilldown:
            base += "\n\n" + drilldown
        return base

    def _format_inventory(self, inv: dict) -> list[str]:
        if not inv:
            return ["\n인벤토리: 비어있음"]
        lines = ["\n## 인벤토리 (핵심 아이템)"]
        for name, count in sorted(inv.items(), key=lambda x: -x[1])[:20]:
            lines.append(f"  {name}: {count}")
        return lines

    def _guess_zone_role(self, z: dict) -> str:
        miners    = z.get("miners", 0)
        furnaces  = z.get("furnaces", 0)
        assemblers= z.get("assemblers", 0)
        turrets   = z.get("turrets", 0)
        labs      = z.get("labs", 0)
        pumps     = z.get("pumps", 0)

        if turrets > 3:          return "방어구역"
        if labs > 0:             return "연구실"
        if pumps > 0:            return "정유/전력"
        if assemblers > miners:  return "조립라인"
        if furnaces > miners:    return "제련구역"
        if miners > 0:           return "채굴구역"
        return "기타"