"""
inference.py — OpenEnv-compliant inference entry point for the
Interior Design Spatial Engine (Furniture Tetris + Human Flow).

STDOUT FORMAT (strictly followed):
    [START] task=<task> env=<benchmark> model=<model>
    [STEP]  step=<n> action=<str> reward=<0.00> done=<true|false> error=<msg|null>
    [END]   success=<true|false> steps=<n> score=<0.000> rewards=<r1,r2,...>
"""

from __future__ import annotations
import sys
import os
import json
import argparse
from typing import List, Optional

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from env import InteriorEnv
from baseline import BaselineAgent
from models import Action
from tasks import TASKS

# ── Environment variables ─────────────────────────────────────────────────────
API_BASE_URL     = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
MODEL_NAME       = os.getenv("MODEL_NAME",   "Qwen/Qwen2.5-72B-Instruct")
HF_TOKEN         = os.getenv("HF_TOKEN")
LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME")
BENCHMARK    = "interior-env"

# ── OpenAI client (optional LLM agent) ───────────────────────────────────────
try:
    from openai import OpenAI
    _llm_client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN) if HF_TOKEN else None
except Exception:
    _llm_client = None

# ── Strict stdout loggers ─────────────────────────────────────────────────────
def log_start(task: str, model: str) -> None:
    print(f"[START] task={task} env={BENCHMARK} model={model}", flush=True)

def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None:
    err = error if error else "null"
    print(f"[STEP] step={step} action={action} reward={reward:.2f} done={str(done).lower()} error={err}", flush=True)

def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:
    rew_str = ",".join(f"{r:.2f}" for r in rewards)
    print(f"[END] success={str(success).lower()} steps={steps} score={score:.3f} rewards={rew_str}", flush=True)

# ── LLM Agent ────────────────────────────────────────────────────────────────
SYSTEM_PROMPT = """You are a spatial AI agent designing interior layouts on a 2D grid.
You place furniture ensuring strict geometrical, human flow, and physics constraints.
- Light rays shoot purely orthogonally from windows. Tall items block them (cast shadows).
- Windows/doors emit cold (-2). Radiators/TVs emit radial heat. 
Output ONLY valid JSON, no explanation:
{"action_type": "place_item", "item": "<item_id>", "x": <int>, "y": <int>, "rotation": <0|90>}"""

def llm_pick_action(obs, step: int) -> Optional[Action]:
    if not _llm_client:
        return None
    try:
        obs_summary = {
            "room_size": obs.room_size,
            "door_position": obs.door_position,
            "windows": obs.windows,
            "unplaced_items": [{"id": u.id, "width": u.width, "height": u.height, "rules": u.rules}
                               for u in obs.unplaced_items],
            "placed_items": {k: {"x": v.x, "y": v.y, "rotation": v.rotation}
                            for k, v in obs.placed_items.items()},
            "feedback": obs.feedback_logs[-3:] if obs.feedback_logs else [],
        }
        completion = _llm_client.chat.completions.create(
            model=MODEL_NAME,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user",   "content": f"Step {step}. Current state:\n{json.dumps(obs_summary, indent=2)}\n\nWhat is your next action?"},
            ],
            response_format={"type": "json_object"},
            temperature=0.2,
            max_tokens=100,
        )
        data = json.loads(completion.choices[0].message.content)
        return Action(
            action_type=data.get("action_type", "no_op"),
            item=data.get("item", ""),
            x=int(data.get("x", 0)),
            y=int(data.get("y", 0)),
            rotation=int(data.get("rotation", 0)),
        )
    except Exception as e:
        return None

# ── Action → compact string ─────────────────────────────────────────────────
def action_str(action: Action) -> str:
    if action.action_type == "place_item":
        return f"place_item({action.item},{action.x},{action.y},{action.rotation})"
    if action.action_type == "remove_item":
        return f"remove_item({action.item})"
    return "no_op()"

# ── Single episode runner ─────────────────────────────────────────────────────
def run_task(task_id: str) -> None:
    if task_id not in TASKS:
        print(f"[ERROR] Unknown task: {task_id}", flush=True)
        return

    config   = TASKS[task_id]
    env      = InteriorEnv()
    baseline = BaselineAgent()

    obs  = env.reset(config, task_id)
    done = False
    step = 0
    rewards: List[float] = []
    success = False
    score   = 0.0

    log_start(task=task_id, model=MODEL_NAME)

    try:
        while not done:
            step += 1

            # Try LLM first, fall back to deterministic baseline
            action = llm_pick_action(obs, step)
            if action is None:
                action = baseline.pick_action(obs, task_id)

            obs, reward, done, _ = env.step(action)
            rewards.append(reward)

            error = None
            if obs.feedback_logs:
                last = obs.feedback_logs[-1]
                if any(tag in last for tag in ("ERROR", "COLLISION", "BOUNDS", "FLOW")):
                    error = last[:80].replace(" ", "_")

            log_step(step=step, action=action_str(action),
                     reward=reward, done=done, error=error)

        result  = env.get_episode_result()
        success = result.success
        score   = result.score

    except Exception as exc:
        log_step(step=step + 1, action="error()", reward=0.0, done=True, error=str(exc)[:80])

    finally:
        log_end(success=success, steps=step, score=score, rewards=rewards)

# ── Entry point ───────────────────────────────────────────────────────────────
def main() -> None:
    parser = argparse.ArgumentParser(description="Interior Env OpenEnv inference")
    parser.add_argument("--task", choices=list(TASKS.keys()),
                        help="Run a single task. Omit to run all levels.")
    args = parser.parse_args()

    tasks_to_run = [args.task] if args.task else list(TASKS.keys())
    for tid in tasks_to_run:
        run_task(tid)

if __name__ == "__main__":
    main()