import json
import random
from pathlib import Path
from typing import Union
from google.adk.tools import FunctionTool   # ← ADK, not genai


def build_adk_tool(
    sops_folder: Union[str, Path],
    dropout_rate: float = 0.0,
) -> FunctionTool:
    sops_folder = Path(sops_folder)
    available_sops = _get_available_sops(sops_folder)

    if dropout_rate > 0 and available_sops:
        keep_count = max(1, int(len(available_sops) * (1 - dropout_rate)))
        available_sops = sorted(
            random.sample(available_sops, keep_count), key=lambda x: x["id"]
        )

    sop_list = "\n".join(
        f'  • {s["id"]} | {_clean_name(s["name"])}: {_clean_desc(s["description"])}'
        for s in available_sops
    )

    # Build lookup cache as a closure so the function can use it
    cache = _build_lookup_cache(sops_folder)

    # The function IS the tool — schema comes from the docstring + type hints
    def read_sop(sop_identifiers: list[str]) -> str:
        results = []
        for identifier in sop_identifiers:
            data = _lookup_sop(cache, str(identifier))
            if data:
                results.append({"identifier": identifier, "status": "success", "data": data})
            else:
                results.append({"identifier": identifier, "status": "not_found",
                                 "error": f"SOP '{identifier}' not found"})
        return json.dumps(results, ensure_ascii=False, indent=2)

    # Dynamically set the description with the embedded SOP list
    read_sop.__doc__ = (
        f"Read Standard Operating Procedures from knowledge base.\n\n"
        f"Available SOPs ({len(available_sops)} total):\n"
        f"{sop_list}\n\n"
        f'Use SOP ID to retrieve (e.g., "290", "1237"). '
        f"You can also use keywords from the name."
    )

    return FunctionTool(func=read_sop)

def _build_lookup_cache(sops_folder: Path) -> dict:
    cache = {}
    for json_file in sops_folder.rglob("*.json"):
        try:
            with open(json_file, "r", encoding="utf-8") as f:
                data = json.load(f)
            if "SOP" in data:
                sop = data["SOP"]
                sop_id = str(sop.get("id", ""))
                folder_name = json_file.parent.name
                name = sop.get("name", "")
                cache[sop_id] = json_file
                cache[name] = json_file
                cache[f"{name} {folder_name.split(' ')[-1]}"] = json_file
        except Exception:
            pass
    return cache


def _lookup_sop(cache: dict, identifier: str):
    if identifier in cache:
        try:
            with open(cache[identifier], "r", encoding="utf-8") as f:
                return json.load(f)
        except Exception:
            return None
    # Fuzzy fallback
    identifier_lower = identifier.lower()
    for key, path in cache.items():
        if identifier_lower in key.lower():
            try:
                with open(path, "r", encoding="utf-8") as f:
                    return json.load(f)
            except Exception:
                continue
    return None

def _get_available_sops(sops_folder: Path) -> list[dict]:
    sops = []
    if not sops_folder.exists():
        return sops
    for json_file in sops_folder.rglob("*.json"):
        try:
            with open(json_file, "r", encoding="utf-8") as f:
                data = json.load(f)
            if "SOP" in data:
                sop = data["SOP"]
                folder_name = json_file.parent.name
                name = sop.get("name", "unknown")
                sops.append({
                    "id": sop.get("id", "unknown"),
                    "name": f"{name} {folder_name.split(' ')[-1]}",
                    "description": sop.get("description", ""),
                })
        except Exception:
            pass
    return sorted(sops, key=lambda x: x["id"])


def _clean_name(name: str) -> str:
    for token in ("＜TECH＞", "［P］", "［H］", "［プ］", "［ス］", "［", "］"):
        name = name.replace(token, "")
    return name.strip()


def _clean_desc(desc: str) -> str:
    return (
        desc.replace("に関する顧客対応手順を記述する。", "")
            .replace("に関する顧客対応手順", "")
            .strip()
    )