from __future__ import annotations

import io
import json
from pathlib import Path

from mrp import Environment


def _transport(*, input=None, files=None, output=None):
    """Build a minimal MRP transport dict."""
    data = {}
    if input is not None:
        data["input"] = input
    if files is not None:
        data.setdefault("model", {})["files"] = files
    if output is not None:
        data["output"] = output
    return data


# --- Constructor ---


class TestInit:
    def test_extracts_input(self):
        ctx = Environment(_transport(input={"r0": 2.5, "gamma": 0.1}))
        assert ctx.input == {"r0": 2.5, "gamma": 0.1}

    def test_pops_replicate(self):
        ctx = Environment(_transport(input={"r0": 2.5, "seed": 42, "replicate": 3}))
        assert ctx.replicate == 3
        assert "replicate" not in ctx.input
        assert ctx.input == {"r0": 2.5, "seed": 42}

    def test_seed_passes_through(self):
        ctx = Environment(_transport(input={"seed": 42, "r0": 2.5}))
        assert ctx.input["seed"] == 42

    def test_defaults_replicate_to_zero(self):
        ctx = Environment(_transport(input={"r0": 2.5}))
        assert ctx.replicate == 0

    def test_coerces_replicate_to_int(self):
        ctx = Environment(_transport(input={"replicate": "2"}))
        assert ctx.replicate == 2

    def test_empty_data(self):
        ctx = Environment({})
        assert ctx.input == {}
        assert ctx.replicate == 0
        assert ctx.files == {}
        assert ctx.output_dir is None

    def test_files_as_paths(self):
        ctx = Environment(
            _transport(
                files={"population": "/data/pop.csv", "geo": "relative/geo.json"}
            )
        )
        assert ctx.files["population"] == Path("/data/pop.csv")
        assert ctx.files["geo"] == Path("relative/geo.json")
        assert isinstance(ctx.files["population"], Path)

    def test_input_is_a_copy(self):
        original = {"r0": 2.5}
        ctx = Environment(_transport(input=original))
        ctx.input["r0"] = 999
        assert original["r0"] == 2.5


# --- output_dir ---


class TestOutputDir:
    def test_filesystem_sink(self):
        ctx = Environment(
            _transport(output={"spec": "filesystem", "dir": "./results/run_0000/"})
        )
        assert ctx.output_dir == Path("./results/run_0000/")

    def test_no_output(self):
        ctx = Environment(_transport())
        assert ctx.output_dir is None

    def test_non_filesystem_output_ignored(self):
        ctx = Environment(
            _transport(output={"spec": "az", "container": "my-container"})
        )
        assert ctx.output_dir is None

    def test_profiled_output_default(self):
        ctx = Environment(
            _transport(
                output={
                    "profile": {
                        "default": {
                            "spec": "filesystem",
                            "dir": "./output/run_0000/",
                        },
                        "stdout": {"spec": "stdout"},
                    }
                }
            )
        )
        assert ctx.output_dir == Path("./output/run_0000/")

    def test_profiled_output_first_key(self):
        ctx = Environment(
            _transport(
                output={
                    "profile": {
                        "local": {
                            "spec": "filesystem",
                            "dir": "./local/run_0000/",
                        },
                    }
                }
            )
        )
        assert ctx.output_dir == Path("./local/run_0000/")


# --- load (stdin) ---


class TestLoadStdin:
    def test_reads_stdin_by_default(self, monkeypatch):
        data = _transport(input={"r0": 3.0, "seed": 7})
        monkeypatch.setattr("sys.stdin", io.StringIO(json.dumps(data)))
        ctx = Environment().load()
        assert ctx.input == {"r0": 3.0, "seed": 7}

    def test_skips_stdin_when_disabled(self, monkeypatch):
        monkeypatch.setattr(
            "sys.stdin", io.StringIO(json.dumps(_transport(input={"r0": 3.0})))
        )
        ctx = Environment().load(stdin=False)
        assert ctx.input == {}

    def test_skips_stdin_when_tty(self, monkeypatch):
        fake = io.StringIO("")
        fake.isatty = lambda: True
        monkeypatch.setattr("sys.stdin", fake)
        ctx = Environment().load()
        assert ctx.input == {}


# --- from_stdin ---


class TestFromStdin:
    def test_parses_json_from_stdin(self, monkeypatch):
        data = _transport(input={"r0": 3.0, "seed": 10})
        monkeypatch.setattr("sys.stdin", io.StringIO(json.dumps(data)))
        ctx = Environment.from_args("stdin")
        assert ctx.input == {"r0": 3.0, "seed": 10}

    def test_empty_stdin_returns_empty(self, monkeypatch):
        monkeypatch.setattr("sys.stdin", io.StringIO(""))
        ctx = Environment.from_args("stdin")
        assert ctx.input == {}

    def test_whitespace_stdin_returns_empty(self, monkeypatch):
        monkeypatch.setattr("sys.stdin", io.StringIO("   \n  "))
        ctx = Environment.from_args("stdin")
        assert ctx.input == {}


# --- from_run_json ---


class TestFromRunJson:
    def test_constructs_from_dict(self):
        data = _transport(input={"r0": 3.0, "seed": 10})
        ctx = Environment.from_args(data)
        assert ctx.input == {"r0": 3.0, "seed": 10}


# --- write ---


class TestWrite:
    def test_writes_string_to_filesystem(self, tmp_path):
        ctx = Environment(
            _transport(output={"spec": "filesystem", "dir": str(tmp_path / "out")})
        )
        ctx.write("hello.txt", "hello world")
        assert (tmp_path / "out" / "hello.txt").read_text() == "hello world"

    def test_writes_bytes_to_filesystem(self, tmp_path):
        ctx = Environment(
            _transport(output={"spec": "filesystem", "dir": str(tmp_path / "out")})
        )
        ctx.write("data.bin", b"\x00\x01\x02")
        assert (tmp_path / "out" / "data.bin").read_bytes() == b"\x00\x01\x02"

    def test_creates_nested_dirs(self, tmp_path):
        deep = tmp_path / "a" / "b" / "c"
        ctx = Environment(_transport(output={"spec": "filesystem", "dir": str(deep)}))
        ctx.write("f.txt", "ok")
        assert (deep / "f.txt").read_text() == "ok"

    def test_string_to_stdout_when_no_sink(self, monkeypatch):
        buf = io.StringIO()
        monkeypatch.setattr("sys.stdout", buf)
        ctx = Environment(_transport())
        ctx.write("ignored.txt", "stdout content")
        assert buf.getvalue() == "stdout content"

    def test_bytes_to_stdout_when_no_sink(self, monkeypatch):
        buf = io.BytesIO()
        monkeypatch.setattr("sys.stdout", type("FakeStdout", (), {"buffer": buf})())
        ctx = Environment(_transport())
        ctx.write("ignored.bin", b"\xff\xfe")
        assert buf.getvalue() == b"\xff\xfe"


# --- write_csv ---


class TestWriteCsv:
    COLUMNS = {"step": [0, 1], "day": [0.0, 1.0], "S": [9990.0, 9985.0]}

    def test_writes_csv_to_filesystem(self, tmp_path):
        ctx = Environment(
            _transport(output={"spec": "filesystem", "dir": str(tmp_path / "out")})
        )
        ctx.write_csv("data.csv", self.COLUMNS)
        content = (tmp_path / "out" / "data.csv").read_text()
        lines = content.strip().split("\n")
        assert lines[0] == "step,day,S"
        assert lines[1] == "0,0.0,9990.0"
        assert lines[2] == "1,1.0,9985.0"

    def test_writes_csv_to_stdout_when_no_sink(self, monkeypatch):
        buf = io.StringIO()
        monkeypatch.setattr("sys.stdout", buf)
        ctx = Environment(_transport())
        ctx.write_csv("ignored.csv", self.COLUMNS)
        lines = buf.getvalue().strip().splitlines()
        assert lines[0] == "step,day,S"
        assert len(lines) == 3

    def test_empty_columns(self, tmp_path):
        ctx = Environment(
            _transport(output={"spec": "filesystem", "dir": str(tmp_path)})
        )
        ctx.write_csv("empty.csv", {"a": [], "b": []})
        content = (tmp_path / "empty.csv").read_text()
        assert content.strip() == "a,b"


# --- csv_writer (streaming) ---


class TestCsvWriter:
    def test_streaming_to_filesystem(self, tmp_path):
        ctx = Environment(
            _transport(output={"spec": "filesystem", "dir": str(tmp_path / "out")})
        )
        with ctx.csv_writer("data.csv", ["step", "value"]) as w:
            w.write_row([0, 1.5])
            w.write_row([1, 2.5])
        content = (tmp_path / "out" / "data.csv").read_text()
        lines = content.strip().split("\n")
        assert lines[0] == "step,value"
        assert lines[1] == "0,1.5"
        assert lines[2] == "1,2.5"

    def test_streaming_dict_rows(self, tmp_path):
        ctx = Environment(
            _transport(output={"spec": "filesystem", "dir": str(tmp_path)})
        )
        with ctx.csv_writer("data.csv", ["a", "b"]) as w:
            w.write_row({"b": 2, "a": 1})
        content = (tmp_path / "data.csv").read_text()
        lines = content.strip().split("\n")
        assert lines[0] == "a,b"
        assert lines[1] == "1,2"

    def test_streaming_to_stdout(self, monkeypatch):
        buf = io.StringIO()
        monkeypatch.setattr("sys.stdout", buf)
        ctx = Environment(_transport())
        w = ctx.csv_writer("ignored.csv", ["x", "y"])
        w.write_row([1, 2])
        w.close()
        lines = buf.getvalue().strip().splitlines()
        assert lines[0] == "x,y"
        assert lines[1] == "1,2"

    def test_headers_only(self, tmp_path):
        ctx = Environment(
            _transport(output={"spec": "filesystem", "dir": str(tmp_path)})
        )
        with ctx.csv_writer("empty.csv", ["a", "b"]) as _w:
            pass
        content = (tmp_path / "empty.csv").read_text()
        assert content.strip() == "a,b"


# --- create_csv / write_csv_row (stateful) ---


class TestStatefulCsv:
    def test_create_and_write_rows(self, tmp_path):
        ctx = Environment(
            _transport(output={"spec": "filesystem", "dir": str(tmp_path)})
        )
        ctx.create_csv("out", "data.csv", ["a", "b"])
        ctx.write_csv_row("out", [1, 2])
        ctx.write_csv_row("out", [3, 4])
        ctx.close_csv("out")
        content = (tmp_path / "data.csv").read_text()
        lines = content.strip().split("\n")
        assert lines[0] == "a,b"
        assert lines[1] == "1,2"
        assert lines[2] == "3,4"

    def test_multiple_writers(self, tmp_path):
        ctx = Environment(
            _transport(output={"spec": "filesystem", "dir": str(tmp_path)})
        )
        ctx.create_csv("a", "a.csv", ["x"])
        ctx.create_csv("b", "b.csv", ["y"])
        ctx.write_csv_row("a", [1])
        ctx.write_csv_row("b", [2])
        ctx.close_all_csv()
        assert (tmp_path / "a.csv").read_text().strip() == "x\n1"
        assert (tmp_path / "b.csv").read_text().strip() == "y\n2"

    def test_dict_rows(self, tmp_path):
        ctx = Environment(
            _transport(output={"spec": "filesystem", "dir": str(tmp_path)})
        )
        ctx.create_csv("out", "data.csv", ["a", "b"])
        ctx.write_csv_row("out", {"b": 2, "a": 1})
        ctx.close_csv("out")
        lines = (tmp_path / "data.csv").read_text().strip().split("\n")
        assert lines[1] == "1,2"
