"""
Gemma 4 E4B Weight-Space Loader & Live Web Search Agent
Loads only the text backbone of `unsloth/gemma-4-E4B-it` in bfloat16:
- Active 42 Transformer layers + tied embed_tokens/lm_head (~9.4 GB) on CUDA (RTX 4070 Ti)
- Static 5.64 GB Per-Layer Embedding table (`embed_tokens_per_layer`) on CPU RAM
- Skips loading vision_tower and audio_tower (~0.95 GB) entirely
"""

import re
import json
import torch
from safetensors import safe_open
from huggingface_hub import hf_hub_download
from transformers import AutoConfig, AutoTokenizer, Gemma4ForCausalLM
from accelerate import init_empty_weights
from ddgs import DDGS


MODEL_REPO = "unsloth/gemma-4-E4B-it"


class CPUOffloadedPLE(torch.nn.Module):
    """Wraps Gemma4TextScaledWordEmbedding on CPU and moves only the small [B, S, 42*256] lookup slice to CUDA."""
    def __init__(self, cpu_embedding: torch.nn.Module, target_device: torch.device):
        super().__init__()
        self.cpu_embedding = cpu_embedding
        self.target_device = target_device

    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        out_cpu = self.cpu_embedding(input_ids.to("cpu"))
        return out_cpu.to(self.target_device, non_blocking=True)


def load_gemma4_e4b_text(repo_id: str = MODEL_REPO, device: str = "cuda"):
    print(f"[Loader] Loading config and tokenizer from {repo_id}...")
    full_config = AutoConfig.from_pretrained(repo_id)
    text_config = full_config.text_config
    text_config.torch_dtype = torch.bfloat16
    tokenizer = AutoTokenizer.from_pretrained(repo_id)

    safetensors_path = hf_hub_download(repo_id=repo_id, filename="model.safetensors")

    print("[Loader] Initializing empty Gemma4ForCausalLM skeleton...")
    with init_empty_weights():
        model = Gemma4ForCausalLM(text_config)

    # Stream tensors from safetensors directly to their target devices (zero RAM spike)
    print(f"[Loader] Streaming weights from {safetensors_path}...")
    target_device = torch.device(device)

    with safe_open(safetensors_path, framework="pt", device="cpu") as f:
        keys = f.keys()
        for k in keys:
            if not k.startswith("model.language_model."):
                # Skip vision_tower, audio_tower, embed_vision, embed_audio
                continue

            # Map key from Gemma4ForConditionalGeneration -> Gemma4ForCausalLM
            # 'model.language_model.X' -> 'model.X'
            sub_key = k.replace("model.language_model.", "model.", 1)

            # Navigate module hierarchy
            parts = sub_key.split(".")
            mod = model
            skip = False
            for p in parts[:-1]:
                if not hasattr(mod, p):
                    skip = True
                    break
                mod = getattr(mod, p)
            if skip or not hasattr(mod, parts[-1]):
                continue

            tensor = f.get_tensor(k)
            # Keep the giant 5.64 GB Per-Layer Embedding table on CPU RAM
            if "embed_tokens_per_layer" in sub_key:
                param_device = "cpu"
            else:
                param_device = target_device

            param = torch.nn.Parameter(tensor.to(device=param_device, dtype=torch.bfloat16), requires_grad=False)
            setattr(mod, parts[-1], param)

    # Initialize any non-persistent buffers (like rotary_emb inv_freq or layer_scalar if not in safetensors)
    for name, buf in model.named_buffers():
        if buf.device.type == "meta":
            parts = name.split(".")
            mod = model
            for p in parts[:-1]:
                mod = getattr(mod, p)
            if parts[-1] == "layer_scalar":
                setattr(mod, parts[-1], torch.ones(1, device=target_device, dtype=torch.bfloat16))

    # Re-init rotary_emb on target_device
    from transformers.models.gemma4.modeling_gemma4 import Gemma4TextRotaryEmbedding
    model.model.rotary_emb = Gemma4TextRotaryEmbedding(text_config).to(target_device)

    # Tie lm_head.weight to model.embed_tokens.weight on GPU
    model.lm_head.weight = model.model.embed_tokens.weight

    # Wrap embed_tokens_per_layer so CPU lookup seamlessly transfers small output slice to GPU
    model.model.embed_tokens_per_layer = CPUOffloadedPLE(model.model.embed_tokens_per_layer, target_device)

    model.eval()
    vram_gb = torch.cuda.memory_allocated(target_device) / 1e9
    print(f"[Loader] Ready! GPU VRAM allocated: {vram_gb:.2f} GB")
    return model, tokenizer


# Standard Web Search Tool Schema for Gemma 4
WEB_SEARCH_TOOL = {
    "type": "function",
    "function": {
        "name": "web_search",
        "description": "Search the live internet for current events, real-time data, recent facts, or specific information you are not 100% certain about.",
        "parameters": {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "The search query to look up on the internet."
                }
            },
            "required": ["query"]
        }
    }
}


def execute_web_search(query: str, max_results: int = 3) -> str:
    """Execute a live DuckDuckGo web search."""
    try:
        results = DDGS().text(query, max_results=max_results)
        if not results:
            return "No results found."
        formatted = []
        for r in results:
            formatted.append(f"Title: {r.get('title', '')}\nURL: {r.get('href', '')}\nSnippet: {r.get('body', '')}")
        return "\n\n".join(formatted)
    except Exception as e:
        return f"Search error: {e}"


def parse_gemma4_tool_call(raw_output: str):
    """
    Parse Gemma 4 native tool call format:
    <|tool_call>call:web_search{query:<|"|>...<|"|>}<tool_call|>
    """
    match = re.search(r"<\|tool_call>call:([a-zA-Z0-9_]+)\{(.*?)\}<tool_call\|>", raw_output, re.DOTALL)
    if not match:
        return None
    func_name = match.group(1)
    args_body = match.group(2)
    args = {}
    for kv in re.finditer(r"([a-zA-Z0-9_]+):<\|\"\|>(.*?)<\|\"\|>", args_body, re.DOTALL):
        args[kv.group(1)] = kv.group(2)
    return {"name": func_name, "arguments": args}