#!/usr/bin/env python3
"""Install the Hermes post-LLM response transform before persistence."""

from __future__ import annotations

import argparse
import ast
import hashlib
import json
import os
import tempfile
from pathlib import Path

PATCH_VERSION = "teleo-response-transform-hook-v1"
PATCH_MARKER = "        # TELEO_RESPONSE_TRANSFORM_HOOK_V1\n"
START_ANCHOR = "        # Save trajectory if enabled\n"
END_ANCHOR = "        # Extract reasoning from the last assistant message (if any)\n"

REPLACEMENT = '''        # TELEO_RESPONSE_TRANSFORM_HOOK_V1
        # Run response validators before trajectory/session persistence. A plugin
        # may return {"assistant_response": "..."}; the last non-empty value wins.
        # This keeps delivered text and durable conversation memory identical.
        if final_response and not interrupted:
            try:
                from hermes_cli.plugins import invoke_hook as _invoke_hook
                _post_results = _invoke_hook(
                    "post_llm_call",
                    session_id=self.session_id,
                    user_message=original_user_message,
                    assistant_response=final_response,
                    conversation_history=list(messages),
                    model=self.model,
                    platform=getattr(self, "platform", None) or "",
                )
                _original_final_response = final_response
                for _post_result in _post_results:
                    if not isinstance(_post_result, dict):
                        continue
                    _candidate = _post_result.get("assistant_response")
                    if isinstance(_candidate, str) and _candidate.strip():
                        final_response = _candidate.strip()
                if final_response != _original_final_response:
                    _updated_assistant_message = False
                    for _message in reversed(messages):
                        if _message.get("role") == "assistant" and not _message.get("tool_calls"):
                            _message["content"] = final_response
                            _updated_assistant_message = True
                            break
                    if not _updated_assistant_message:
                        messages.append({"role": "assistant", "content": final_response})
            except Exception as exc:
                logger.warning("post_llm_call hook failed: %s", exc)

        # Save trajectory if enabled
        self._save_trajectory(messages, user_message, completed)

        # Clean up VM and browser for this task after conversation completes
        self._cleanup_task_resources(effective_task_id)

        # Persist session to both JSON log and SQLite
        self._persist_session(messages, conversation_history)


'''


def sha256_text(value: str) -> str:
    return hashlib.sha256(value.encode("utf-8")).hexdigest()


def patch_source(source: str) -> tuple[str, bool]:
    if PATCH_MARKER in source:
        required = (
            '_candidate = _post_result.get("assistant_response")',
            "self._save_trajectory(messages, user_message, completed)",
            "self._persist_session(messages, conversation_history)",
        )
        missing = [item for item in required if item not in source]
        if missing:
            raise ValueError(f"installed patch marker is incomplete: {missing}")
        return source, False

    if source.count(START_ANCHOR) != 1 or source.count(END_ANCHOR) != 1:
        raise ValueError("Hermes response lifecycle anchors changed; refusing an unreviewed patch")
    start = source.index(START_ANCHOR)
    end = source.index(END_ANCHOR, start)
    original_block = source[start:end]
    required_original = (
        '"post_llm_call"',
        "self._save_trajectory(messages, user_message, completed)",
        "self._cleanup_task_resources(effective_task_id)",
        "self._persist_session(messages, conversation_history)",
    )
    missing = [item for item in required_original if item not in original_block]
    if missing:
        raise ValueError(f"Hermes response lifecycle changed; missing expected statements: {missing}")

    patched = source[:start] + REPLACEMENT + source[end:]
    ast.parse(patched)
    return patched, True


def install(path: Path, *, check: bool) -> tuple[dict[str, str | bool], int]:
    source = path.read_text(encoding="utf-8")
    patched, changed = patch_source(source)
    before_sha = sha256_text(source)
    after_sha = sha256_text(patched)
    if check:
        status = "patch_required" if changed else "installed"
        return {
            "patch_version": PATCH_VERSION,
            "status": status,
            "target": str(path),
            "sha256": before_sha,
        }, 1 if changed else 0

    if changed:
        mode = path.stat().st_mode & 0o777
        with tempfile.NamedTemporaryFile(
            "w", encoding="utf-8", dir=path.parent, prefix=f".{path.name}.", delete=False
        ) as handle:
            handle.write(patched)
            temp_path = Path(handle.name)
        try:
            temp_path.chmod(mode)
            os.replace(temp_path, path)
        finally:
            temp_path.unlink(missing_ok=True)

    ast.parse(path.read_text(encoding="utf-8"))
    return {
        "patch_version": PATCH_VERSION,
        "status": "installed_now" if changed else "already_installed",
        "target": str(path),
        "before_sha256": before_sha,
        "after_sha256": after_sha,
    }, 0


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("target", type=Path)
    parser.add_argument("--check", action="store_true")
    args = parser.parse_args()
    result, code = install(args.target, check=args.check)
    print(json.dumps(result, sort_keys=True))
    return code


if __name__ == "__main__":
    raise SystemExit(main())