#!/usr/bin/env python3
"""Patch hermes-agent gateway to auto-inject MEDIA: tags from write_file tool calls.

Problem: When the LLM creates files via write_file but doesn't include MEDIA: tags
in its final response, the gateway has no way to detect and deliver those files as
native attachments. The LLM may say "I've sent the file" when it actually only saved
it locally.

Solution: After the existing TTS MEDIA: propagation block in _run_agent(), scan the
assistant messages for write_file tool calls. Extract file paths from the tool call
arguments and append MEDIA: tags for files that:
  1. Actually exist on disk
  2. Have a document/media extension
  3. Are not already referenced in the final_response

This mirrors the existing TTS propagation pattern (gateway/run.py ~line 10968).
"""

import re
import sys
import os
import glob
import json


def patch_gateway(filepath: str):
    with open(filepath, 'r') as f:
        content = f.read()

    # The patch insertion point: right after the TTS MEDIA: propagation block
    # ends with "final_response = final_response + ..."
    old = '''                    final_response = final_response + "\n" + "\n".join(unique_tags)
            
            # Sync session_id: the agent may have created a new session during'''

    new = '''                    final_response = final_response + "\n" + "\n".join(unique_tags)
            
            # Auto-inject MEDIA: tags for write_file tool results.
            # When the LLM creates files via write_file but forgets to include
            # MEDIA: tags in its response, this ensures the files are still
            # delivered as native attachments on Feishu/WeChat/etc.
            _doc_exts = {
                '.md', '.txt', '.csv', '.json', '.xml', '.yaml', '.yml', '.toml', '.log',
                '.pdf', '.doc', '.docx', '.xls', '.xlsx', '.ppt', '.pptx',
                '.html', '.htm', '.zip', '.tar', '.gz', '.7z', '.rar',
                '.png', '.jpg', '.jpeg', '.gif', '.webp', '.bmp', '.svg',
                '.mp4', '.mov', '.avi', '.mkv', '.webm',
                '.ogg', '.opus', '.mp3', '.wav', '.m4a',
            }
            _auto_media_paths = []
            for _msg in result.get("messages", []):
                if _msg.get("role") != "assistant":
                    continue
                _tool_calls = _msg.get("tool_calls") or []
                for _tc in _tool_calls:
                    _fn = (_tc.get("function") or {})
                    if not isinstance(_fn, dict):
                        continue
                    _fn_name = _fn.get("name", "")
                    if _fn_name not in ("write_file",):
                        continue
                    try:
                        _args = json.loads(_fn.get("arguments", "{}"))
                    except (json.JSONDecodeError, TypeError):
                        continue
                    _fpath = _args.get("path", "")
                    if not _fpath:
                        continue
                    _fpath = os.path.expanduser(_fpath)
                    # Only inject for files with document/media extensions
                    _ext = os.path.splitext(_fpath)[1].lower()
                    if _ext not in _doc_exts:
                        continue
                    # File must actually exist
                    if not os.path.isfile(_fpath):
                        continue
                    # Avoid duplicates and paths already in response
                    _media_tag = f"MEDIA:{_fpath}"
                    if _media_tag in final_response:
                        continue
                    if _fpath in final_response:
                        continue
                    _auto_media_paths.append(_media_tag)

            if _auto_media_paths:
                logger.info(
                    "Auto-injecting %d MEDIA: tag(s) from write_file tool calls: %s",
                    len(_auto_media_paths),
                    ", ".join(os.path.basename(p.split(":")[1]) for p in _auto_media_paths),
                )
                final_response = final_response + "\n" + "\n".join(_auto_media_paths)
            
            # Sync session_id: the agent may have created a new session during'''

    if old not in content:
        print(f"WARNING: Could not find insertion point in {filepath}", file=sys.stderr)
        print("The TTS MEDIA: propagation block may have changed. Skipping this patch.", file=sys.stderr)
        sys.exit(0)

    content = content.replace(old, new, 1)

    with open(filepath, 'w') as f:
        f.write(content)

    print(f"Patched {filepath}: auto-inject MEDIA: tags from write_file tool calls")


if __name__ == "__main__":
    # hermes-agent is installed in editable mode (-e), so source is in /app/hermes-agent/
    candidates = [
        "/app/hermes-agent/gateway/run.py",
    ]

    # Also search venv site-packages as fallback
    candidates.extend(glob.glob("/app/venv/lib/**/gateway/run.py", recursive=True))

    filepath = None
    for c in candidates:
        if os.path.isfile(c):
            filepath = c
            break

    if not filepath:
        print("WARNING: run.py not found in any candidate location", file=sys.stderr)
        print(f"Checked: {candidates}", file=sys.stderr)
        print("Skipping patch_auto_media.", file=sys.stderr)
        sys.exit(0)

    patch_gateway(filepath)