#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
JavaScript code generator for CustomProtocolPlugin.

Generates Frida JS from hook definitions:
  - Module watcher matching library_patterns
  - Function resolution (symbol -> pattern -> offset fallback)
  - Interceptor.attach() for each hook
  - sendKey() / sendData() helper functions
  - Key extraction (on_enter / on_leave)
  - Data capture hooks (read/write)
  - Custom JS snippets
"""

from __future__ import annotations

import json
from functools import lru_cache
from pathlib import Path
from typing import Any, Dict, List, Optional

_SHARED_DIR = Path(__file__).parent / "shared_utility"


@lru_cache(maxsize=None)
def _read_js(filename: str) -> str:
    """Read a JS file from the shared_utility directory (cached)."""
    return (_SHARED_DIR / filename).read_text(encoding="utf-8")


def generate_js(
    hook_defs: List[Dict[str, Any]],
    library_patterns: List[str],
    protocol_name: str,
    platforms: Optional[List[str]] = None,
    architectures: Optional[List[str]] = None,
    debug: bool = False,
) -> str:
    """Generate complete Frida JS source from hook definitions.

    Parameters
    ----------
    hook_defs : list
        Serialized HookDefinition dicts from CustomProtocolPlugin.
    library_patterns : list
        Regex patterns for target library detection.
    protocol_name : str
        Protocol identifier for message tagging.
    platforms : list, optional
        Restrict to these platforms (empty = all).
    architectures : list, optional
        Restrict to these architectures (empty = all).
    debug : bool
        Enable verbose devlog output.

    Returns
    -------
    str
        Complete Frida JavaScript source string.
    """
    platforms = platforms or []
    architectures = architectures or []

    parts: List[str] = []
    parts.append(_generate_header(protocol_name, debug))
    parts.append(_generate_helpers())
    parts.append(_generate_resolve_function())
    parts.append(_generate_hooks(hook_defs))
    parts.append(_generate_module_watcher(library_patterns, platforms, architectures))
    return "\n".join(parts)


def _generate_header(protocol_name: str, debug: bool) -> str:
    """Generate the JS file header with globals."""
    return f"""\
// Auto-generated by friTap CustomProtocolPlugin
// Protocol: {protocol_name}
'use strict';

var _protocol = {json.dumps(protocol_name)};
var _debug = {json.dumps(debug)};
var _hooked_modules = {{}};
"""


def _generate_helpers() -> str:
    """Generate JS helper functions available in all hook callbacks."""
    return _read_js("helpers.js") + "\n" + _read_js("messaging.js")


def _generate_resolve_function() -> str:
    """Generate the function resolution helper (symbol -> pattern -> offset)."""
    return _read_js("resolve.js")


def _generate_hooks(hook_defs: List[Dict[str, Any]]) -> str:
    """Generate the installHooks() function with all hook definitions."""
    parts: List[str] = []
    parts.append("// --- Hook definitions ---")
    parts.append("")
    parts.append(f"var HOOK_DEFS = {json.dumps(hook_defs, indent=2)};")
    parts.append("")
    parts.append(_generate_install_hooks())
    return "\n".join(parts)


def _generate_byte_offset_js(hook_def: Dict[str, Any], var_name: str = "byteOffset") -> str:
    """Generate JS for byte_offset that may be int or per-arch dict."""
    byte_offset = hook_def.get("byte_offset", 0)
    if isinstance(byte_offset, dict):
        # Per-architecture byte offset
        lines = [f"    var {var_name};"]
        lines.append("    switch (Process.arch) {")
        for arch, off in byte_offset.items():
            lines.append(f'        case "{arch}": {var_name} = {off}; break;')
        lines.append(f"        default: {var_name} = 0; break;")
        lines.append("    }")
        return "\n".join(lines)
    else:
        return f"    var {var_name} = {byte_offset};"


def _generate_install_hooks() -> str:
    """Generate the installHooks function that processes HOOK_DEFS."""
    return _read_js("hooks.js")


def _generate_module_watcher(
    library_patterns: List[str],
    platforms: List[str],
    architectures: List[str],
) -> str:
    """Generate the module watcher that triggers hook installation."""
    template = _read_js("module_watcher.js")
    return (
        template
        .replace("%%LIBRARY_PATTERNS%%", json.dumps(library_patterns))
        .replace("%%ALLOWED_PLATFORMS%%", json.dumps(platforms))
        .replace("%%ALLOWED_ARCHITECTURES%%", json.dumps(architectures))
    )
