#!/usr/bin/env python3
"""Shared helpers for the risc-v-gsp-rm-debugger automation harness.

The harness scripts run from the tool root without installation, so this
module bootstraps the tool root onto ``sys.path`` before importing greeneye.
The QEMU backend is the only execution backend the harness claims: the matrix
must stay comparable to the CLI ``probe``/``demo`` paths and must not depend
on the in-development native simulator.
"""
from __future__ import annotations

import contextlib
import hashlib
import json
import os
import shutil
import sys
import tempfile
from dataclasses import dataclass, field

TOOL_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if TOOL_ROOT not in sys.path:
    sys.path.insert(0, TOOL_ROOT)

from greeneye.analyze import parse_qemu_faults  # noqa: E402
from greeneye.elf import FirmwareImage, load_elf, write_flat_kernel  # noqa: E402
from greeneye.errors import GreenEyeError, TargetError  # noqa: E402
from greeneye.rsp import RspClient  # noqa: E402
from greeneye.session import Session  # noqa: E402
from greeneye.symbols import SymbolDB  # noqa: E402
from greeneye.target import QemuTarget, find_qemu  # noqa: E402

DEMO_ELF = os.path.join(TOOL_ROOT, "firmware", "uproc_demo", "build", "uproc_demo.elf")
DEMO_SCHEMA = os.path.join(TOOL_ROOT, "firmware", "uproc_demo", "gsp_ctrl.schema.json")

#: Raw images are relocated to the top of QEMU's virt RAM, matching `greeneye probe`.
LOAD_ADDR = 0x80000000
#: Refuse images whose flat form would be unreasonably large.
MAX_FLAT_BYTES = 256 * 1024 * 1024
#: Cap the QEMU memory size the harness will request.
MAX_MEMORY_MB = 2048
#: Default bound when reading a QEMU log window (trap loops flood `-d int`).
MAX_LOG_BYTES = 1 << 20


def sha256_file(path: str) -> str:
    digest = hashlib.sha256()
    with open(path, "rb") as fh:
        for chunk in iter(lambda: fh.read(1 << 20), b""):
            digest.update(chunk)
    return digest.hexdigest()


def qemu_available() -> bool:
    try:
        find_qemu()
        return True
    except TargetError:
        return False


def qemu_path() -> str:
    return find_qemu()


def load_target(path: str) -> FirmwareImage:
    if not os.path.isfile(path):
        raise GreenEyeError(f"target not found: {path}")
    image = load_elf(path)
    if not image.is_riscv:
        raise GreenEyeError(f"{path}: not a RISC-V ELF image")
    span = image.max_vaddr - image.min_vaddr
    if span > MAX_FLAT_BYTES:
        raise GreenEyeError(f"{path}: image span {span:#x} exceeds the {MAX_FLAT_BYTES:#x} harness cap")
    return image


def memory_mb_for(image: FirmwareImage) -> int:
    span = image.max_vaddr - image.min_vaddr
    return min(MAX_MEMORY_MB, max(256, span // (1024 * 1024) + 64))


def ensure_dir(path: str) -> str:
    os.makedirs(path, exist_ok=True)
    return path


def write_json(path: str, payload: dict) -> str:
    ensure_dir(os.path.dirname(path))
    with open(path, "w", encoding="utf-8") as fh:
        json.dump(payload, fh, indent=2, sort_keys=True)
        fh.write("\n")
    return path


def write_text(path: str, text: str) -> str:
    ensure_dir(os.path.dirname(path))
    with open(path, "w", encoding="utf-8") as fh:
        fh.write(text)
    return path


def relpath(path: str, root: str) -> str:
    return os.path.relpath(path, root)


def counter_delta(a: dict, b: dict) -> list[dict]:
    """Stable union delta of two count maps (B minus A), sorted by key."""
    rows = []
    for key in sorted(set(a) | set(b)):
        a_count = int(a.get(key, 0))
        b_count = int(b.get(key, 0))
        rows.append({"key": key, "a": a_count, "b": b_count, "delta": b_count - a_count})
    return rows


def read_log_window(path: str | None, offset: int = 0, max_bytes: int = MAX_LOG_BYTES) -> tuple[str, bool]:
    """Read a bounded QEMU log window; returns ``(text, capped)``.

    Trap loops emit a ``-d int`` line per faulting instruction, so an
    unbounded read can be hundreds of megabytes. Faults of interest (the
    first exception) are at the start of the window.
    """
    if not path:
        return "", False
    try:
        with open(path, "rb") as fh:
            fh.seek(offset)
            data = fh.read(max_bytes)
    except (FileNotFoundError, OSError):
        return "", False
    capped = len(data) == max_bytes
    text = data.decode("utf-8", "replace")
    if capped:
        text = text.rsplit("\n", 1)[0]
    return text, capped


def read_faults_window(path: str | None, offset: int = 0, max_bytes: int = MAX_LOG_BYTES) -> tuple[list[dict], bool]:
    text, capped = read_log_window(path, offset=offset, max_bytes=max_bytes)
    return parse_qemu_faults(text), capped


@dataclass
class FlatSession:
    """A live QEMU session around a relocated flat kernel."""

    image: FirmwareImage
    bias: int
    session: Session
    target: QemuTarget
    workdir: str
    flat_kernel: str
    owns_workdir: bool = field(default=False)


@contextlib.contextmanager
def open_flat_session(
    image: FirmwareImage,
    target_path: str,
    *,
    record: str | bool | None,
    name: str,
    symbols: SymbolDB | None = None,
    log_faults: bool = True,
    workdir: str | None = None,
    meta: dict | None = None,
):
    """Start QEMU on the target's flat relocated image and yield a session.

    Relocation (``LOAD_ADDR - image.min_vaddr``) mirrors ``greeneye probe``:
    firmware linked for GPU physical addresses has to be placed into virt RAM
    before QEMU can execute it. Symbols are loaded at the relocated addresses.
    """
    owns_workdir = workdir is None
    if workdir is None:
        workdir = tempfile.mkdtemp(prefix="greeneye-harness-")
    flat_kernel = os.path.join(workdir, "flat.bin")
    write_flat_kernel(target_path, flat_kernel, load_addr=LOAD_ADDR)
    bias = LOAD_ADDR - image.min_vaddr
    target = QemuTarget(
        kernel=flat_kernel,
        memory_mb=memory_mb_for(image),
        dlog="guest_errors,unimp,int" if log_faults else None,
        workdir=workdir,
    )
    target.start(paused=True)
    client = RspClient(target.host, target.port)
    client.connect()
    client.sync()
    db = symbols if symbols is not None else SymbolDB()
    session_meta = {"backend": "qemu", "bias": bias, "firmware": target_path, "harness": name}
    session_meta.update(meta or {})
    session = Session(
        client=client,
        target=target,
        symbols=db,
        record=record,
        name=name,
        meta=session_meta,
    )
    session.load_image(image, bias=bias, load_symbols=symbols is None)
    try:
        yield FlatSession(
            image=image,
            bias=bias,
            session=session,
            target=target,
            workdir=workdir,
            flat_kernel=flat_kernel,
            owns_workdir=owns_workdir,
        )
    finally:
        try:
            session.close()
        finally:
            if owns_workdir:
                shutil.rmtree(workdir, ignore_errors=True)