import contextlib
import copy
import os
import unittest
from unittest import mock

from terrarium import object_compiler

from terrarium.object_compiler import (
    DEFAULT_OBJECT_COMPILER_MODEL,
    OBJECT_COMPILER_PROVIDER_TARGET,
    PROVIDER_TARGET,
    ObjectValidationError,
    OpenRouterSettings,
    build_openrouter_payload,
    compile_object,
    deterministic_compile,
    validate_object_payload,
)
from terrarium.sim import TerrariumSim

PROTOTYPE_MODEL = "qwen/qwen3.6-35b-a3b:nitro"

# Explicitly setting these keys (rather than deleting them) keeps load_local_env
# from pulling values out of a real .env.local: it skips keys already in os.environ.
OBJECT_COMPILER_ENV_KEYS = (
    "OBJECT_COMPILER_MODEL",
    "OBJECT_COMPILER_PROVIDER_ONLY",
    "OBJECT_COMPILER_PROVIDER_ORDER",
    "OBJECT_COMPILER_QUANTIZATION",
    "OBJECT_COMPILER_ALLOW_FALLBACKS",
    "OPENROUTER_API_KEY",
    "OPENROUTER_BASE_URL",
    "OPENROUTER_QUANTIZATION",
)


@contextlib.contextmanager
def _patched_env(**overrides):
    # Stub load_local_env so the env is exactly what we set here, then apply only
    # the explicit overrides (removing the rest) so .get(..., default) fallbacks
    # fire and the test never depends on a real .env.local.
    with mock.patch.object(object_compiler, "load_local_env", lambda *a, **k: None):
        with mock.patch.dict(os.environ, overrides, clear=False):
            for key in OBJECT_COMPILER_ENV_KEYS:
                if key not in overrides:
                    os.environ.pop(key, None)
            yield


class ObjectCompilerTest(unittest.TestCase):
    def test_deterministic_compiler_produces_distinct_world_truth(self):
        hot_fruit = deterministic_compile("hot sweet fruit that hurts a little", (8, 8))
        cold_stone = deterministic_compile("cold safe stone", (9, 8))
        humming = deterministic_compile("shiny humming pebble", (10, 8))

        self.assertTrue(hot_fruit["affordances"]["edible"])
        self.assertGreater(hot_fruit["emissions"]["heat"], 0.5)
        self.assertGreater(cold_stone["emissions"]["cold"], 0.5)
        self.assertGreater(cold_stone["emissions"]["home_smell"], 0.4)
        self.assertGreater(humming["emissions"]["sound"], 0.2)
        self.assertEqual(hot_fruit["provenance"]["provider"], "deterministic")

    def test_validation_rejects_interpretive_or_direct_control_fields(self):
        payload = deterministic_compile("sweet fruit", (8, 8))
        invalid = copy.deepcopy(payload)
        invalid["novelty_score"] = 1.0
        with self.assertRaises(ObjectValidationError):
            validate_object_payload(invalid, (8, 8))

    def test_openrouter_payload_uses_required_model_provider_and_schema(self):
        settings = OpenRouterSettings(
            api_key="secret",
            base_url="https://openrouter.ai/api/v1",
            model=PROTOTYPE_MODEL,
            provider_only=PROVIDER_TARGET,
            provider_order=PROVIDER_TARGET,
            require_parameters=True,
            allow_fallbacks=False,
            quantization="fp8",
            structured_outputs_required=True,
        )
        payload = build_openrouter_payload("hot sweet fruit", (8, 8), settings)
        self.assertEqual(payload["model"], PROTOTYPE_MODEL)
        self.assertEqual(payload["provider"]["only"], [PROVIDER_TARGET])
        self.assertFalse(payload["provider"]["allow_fallbacks"])
        self.assertEqual(payload["response_format"]["type"], "json_schema")
        self.assertTrue(payload["response_format"]["json_schema"]["strict"])

    def test_for_object_compiler_reads_object_compiler_env(self):
        with _patched_env(
            OBJECT_COMPILER_MODEL="qwen/qwen3.5-9b",
            OBJECT_COMPILER_PROVIDER_ONLY="deepinfra",
            OBJECT_COMPILER_PROVIDER_ORDER="deepinfra,together",
            OPENROUTER_API_KEY="shared-key",
            OPENROUTER_BASE_URL="https://example/api/v1/",
        ):
            settings = OpenRouterSettings.for_object_compiler()
        self.assertEqual(settings.model, "qwen/qwen3.5-9b")
        self.assertEqual(settings.provider_only_list(), ["deepinfra"])
        self.assertEqual(settings.provider_order_list(), ["deepinfra", "together"])
        # API key and base URL are always shared with OPENROUTER_*.
        self.assertEqual(settings.api_key, "shared-key")
        self.assertEqual(settings.base_url, "https://example/api/v1")

    def test_for_object_compiler_falls_back_to_defaults(self):
        with _patched_env(OPENROUTER_API_KEY="shared-key"):
            settings = OpenRouterSettings.for_object_compiler()
        self.assertEqual(settings.model, DEFAULT_OBJECT_COMPILER_MODEL)
        # Small fast model pinned to its benchmarked-fastest provider path.
        self.assertEqual(settings.provider_order, OBJECT_COMPILER_PROVIDER_TARGET)
        self.assertEqual(settings.quantization, "fp8")
        self.assertTrue(settings.allow_fallbacks)

    def test_compile_object_returns_deterministic_payload_without_provider(self):
        with _patched_env():
            payload = compile_object("a small warm glowing stone", (5, 5))
        self.assertEqual(payload["provenance"]["provider"], "deterministic")
        self.assertEqual(payload["position"], [5, 5])

    def test_create_object_from_text_places_near_user(self):
        sim = TerrariumSim(seed=7)
        state = sim.create_object_from_text("a cold glowing safe stone", provider="deterministic")
        custom = [obj for obj in state["objects"] if obj["id"].startswith("obj_custom")]
        self.assertEqual(len(custom), 1)
        user = state["user"]["position"]
        position = custom[0]["position"]
        self.assertLessEqual(abs(position[0] - user[0]) + abs(position[1] - user[1]), 6)
        self.assertIn("object_placed", [event["type"] for event in state["recent_events"]])

    def test_compiled_object_can_be_placed_and_affects_fields(self):
        sim = TerrariumSim(seed=7)
        payload = deterministic_compile("hot sweet fruit", (8, 8))
        state = sim.place_compiled_object(payload)
        custom = [obj for obj in state["objects"] if obj["id"].startswith("obj_custom")]
        self.assertEqual(len(custom), 1)
        self.assertGreater(state["fields"][8][8]["food_smell"], 0.5)
        self.assertGreater(state["fields"][8][8]["heat"], 0.5)
        self.assertIn("object_placed", [event["type"] for event in state["recent_events"]])


if __name__ == "__main__":
    unittest.main()