"""Unit tests for ELF parsing, symbol resolution and structure decoding."""
from __future__ import annotations

import json
import os
import struct
import tempfile
import unittest

from greeneye.elf import EM_RISCV, parse_elf, load_elf, write_relocated_elf
from greeneye.errors import GreenEyeError
from greeneye.structs import Field, Schema, decode, render, schema_from_json
from greeneye.symbols import Symbol, SymbolDB

DEMO_ELF = os.path.join(os.path.dirname(__file__), "..", "firmware", "uproc_demo", "build", "uproc_demo.elf")


class ElfTests(unittest.TestCase):
    def test_parse_demo_elf(self):
        if not os.path.exists(DEMO_ELF):
            self.skipTest("demo firmware not built")
        image = load_elf(DEMO_ELF)
        self.assertTrue(image.is_riscv)
        self.assertEqual(image.e_machine, EM_RISCV)
        self.assertEqual(image.entry, 0x80000000)
        self.assertGreaterEqual(len(image.load_segments), 1)
        names = {s.name for s in image.symbols}
        self.assertIn("main", names)
        self.assertIn("rpc_ping", names)
        self.assertIn("gsp_ctrl", names)
        ctrl = next(s for s in image.symbols if s.name == "gsp_ctrl")
        self.assertGreaterEqual(ctrl.value, 0x80000000)
        self.assertEqual(ctrl.size, 0x190)
        plan = image.flat_load_plan()
        self.assertTrue(plan and plan[0][0] == 0x80000000)

    def test_not_elf(self):
        with self.assertRaises(Exception):
            parse_elf(b"not an elf")

    def test_write_relocated_elf(self):
        if not os.path.exists(DEMO_ELF):
            self.skipTest("demo firmware not built")
        image = load_elf(DEMO_ELF)
        with tempfile.TemporaryDirectory(prefix="greeneye-reloc-") as tmp:
            out = os.path.join(tmp, "relocated.elf")
            path = write_relocated_elf(DEMO_ELF, out, 0x1000)
            with open(path, "rb") as fh:
                data = fh.read()
        self.assertEqual(struct.unpack_from("<Q", data, 24)[0], image.entry + 0x1000)
        moved = parse_elf(data)
        self.assertEqual(moved.entry, image.entry + 0x1000)


class SymbolTests(unittest.TestCase):
    def setUp(self):
        self.db = SymbolDB()
        self.db.add(Symbol(addr=0x1000, name="alpha", size=0x100, source="test"))
        self.db.add(Symbol(addr=0x1200, name="beta", size=0x20, source="test"))

    def test_resolve(self):
        self.assertEqual(self.db.resolve(0x1010).name, "alpha")
        self.assertEqual(self.db.resolve(0x1210).name, "beta")
        self.assertIsNone(self.db.resolve(0x5000))
        self.assertEqual(self.db.annotate(0x1010), "0x1010 alpha+0x10")
        self.assertEqual(self.db.lookup("beta"), 0x1200)

    def test_by_subsystem(self):
        self.db.add(Symbol(addr=0x2000, name="gamma", size=8, source="test", subsystem="core"))
        self.assertEqual([s.name for s in self.db.by_subsystem("core")], ["gamma"])
        self.assertEqual(self.db.by_subsystem("nope"), [])

    def test_names_tsv(self):
        with tempfile.NamedTemporaryFile("w", suffix=".tsv", delete=False) as fh:
            fh.write("# comment\n0x00279000\t_DT_INIT\tghidra\n0x00279020\tFUN_x\n")
            path = fh.name
        try:
            count = self.db.load_names_tsv(path, bias=0x100000)
            self.assertEqual(count, 2)
            self.assertEqual(self.db.lookup("_DT_INIT"), 0x279000 + 0x100000)
        finally:
            os.unlink(path)

    def test_index_tsv(self):
        with tempfile.NamedTemporaryFile("w", suffix=".tsv", delete=False) as fh:
            fh.write("addr\tfile_addr\tsize\tname\n")
            fh.write("0x00279000\t0x00179000\t24\t_DT_INIT\n")
            path = fh.name
        try:
            count = self.db.load_index_tsv(path, bias=0)
            self.assertEqual(count, 1)
            sym = self.db.get("_DT_INIT")
            self.assertEqual(sym.addr, 0x179000)
            self.assertEqual(sym.size, 24)
            self.assertEqual(sym.extras["program_addr"], 0x279000)
        finally:
            os.unlink(path)

    def test_json_roundtrip(self):
        with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False) as fh:
            path = fh.name
        try:
            self.db.save_json(path, meta={"unit": "test"})
            other = SymbolDB()
            self.assertEqual(other.load_json(path), 2)
            self.assertEqual(other.lookup("alpha"), 0x1000)
            payload = json.load(open(path))
            self.assertEqual(payload["format"], "greeneye-symbols/1")
        finally:
            os.unlink(path)

    def test_json_list_payload(self):
        with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False) as fh:
            json.dump([{"addr": "0x2000", "name": "listed", "size": 4}], fh)
            path = fh.name
        try:
            self.assertEqual(self.db.load_json(path), 1)
            self.assertEqual(self.db.lookup("listed"), 0x2000)
        finally:
            os.unlink(path)

    def test_json_bad_input_raises_greeneye_error(self):
        with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False) as fh:
            fh.write("not json")
            path = fh.name
        try:
            with self.assertRaises(GreenEyeError):
                self.db.load_json(path)
        finally:
            os.unlink(path)


class StructTests(unittest.TestCase):
    SCHEMA = {
        "name": "demo",
        "fields": [
            {"name": "magic", "offset": 0, "fmt": "u32"},
            {"name": "state", "offset": 4, "fmt": "u32", "enum": {"0": "IDLE", "1": "RUN"}},
            {"name": "count", "offset": 8, "fmt": "u64"},
            {"name": "name", "offset": 16, "fmt": "cstr", "size": 16},
            {"name": "buf", "offset": 32, "fmt": "bytes", "size": 8},
        ],
    }

    def test_decode_render(self):
        schema = schema_from_json(self.SCHEMA)
        self.assertEqual(schema.size, 40)
        data = bytearray(40)
        data[0:4] = (0x47535031).to_bytes(4, "little")
        data[4:8] = (1).to_bytes(4, "little")
        data[8:16] = (42).to_bytes(8, "little")
        data[16:22] = b"green\x00"
        data[32:40] = bytes(range(8))
        values = decode(schema, bytes(data), base=0x80000000)
        self.assertEqual(values["magic"], 0x47535031)
        self.assertEqual(values["state"], 1)
        self.assertEqual(values["count"], 42)
        self.assertEqual(values["name"], "green")
        self.assertEqual(values["buf"], bytes(range(8)))
        text = render(schema, values)
        self.assertIn("magic", text)
        self.assertIn("0x47535031", text)
        self.assertIn("green", text)

    def test_invalid_schema(self):
        with self.assertRaises(Exception):
            Field(name="x", offset=0, fmt="nope").nbytes


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