# get_kb.py
import json
import os
from typing import Dict, Any
from agents import function_tool, RunContextWrapper

# KB storage path
KB_PATH = "/tmp/kb.json"

# -------------------------------
# NEW FULL KNOWLEDGE BASE
# -------------------------------
DEFAULT_KB: Dict[str, Any] = {

    "US_PARENT": {
        "description": (
            "The US Parent is the global headquarters and the owner of all intellectual property (IP). "
            "It oversees central management, strategy, risk, and group-wide coordination. "
            "Because it is pre-break-even, it currently operates at a loss and requires funding from other group entities."
        ),
        "receives": (
            "The US entity receives royalties from Luxembourg for the use of US-owned IP. "
            "It also receives dividends from Korea when Korea generates profit, as well as interest, "
            "shareholder loan repayments, or capital contributions when liquidity needs arise."
        ),
        "strategic_role": (
            "The US Parent functions as the strategic hub of the group, consolidating losses and central expenses."
        )
    },

    "KOREA_SUB": {
        "description": (
            "Korea is the group's main engineering and R&D center and holds the majority of the skilled workforce. "
            "It operates under a contract R&D model, providing development services to the US on a cost-plus basis."
        ),
        "liquidity": (
            "Korea holds approximately 90% of the group’s cash, making it the primary source of liquidity for operations."
        ),
        "financial_position": (
            "Korea is currently in cash surplus and is expected to provide liquidity support to the US and Luxembourg "
            "through shareholder loans or equity injections."
        ),
        "strategic_role": (
            "Korea acts as the operational backbone and main liquidity reservoir of the group."
        )
    },

    "LUXEMBOURG_SUB": {
        "description": (
            "Luxembourg serves as the European Regional Headquarters and is the principal contracting entity for "
            "all EU customers. It signs customer contracts, books all EU revenue, and manages regional commercial activity."
        ),
        "financial_position": (
            "Luxembourg is currently loss-making but is projected to become highly profitable within three years. "
            "Once profitable, it will operate as the residual profit center for the EU market."
        ),
        "payments": (
            "Luxembourg pays royalties to the US for the use of IP and pays service fees to France and Korea "
            "for sales support and R&D services, respectively."
        ),
        "strategic_role": (
            "Luxembourg acts as the commercial principal for Europe, taking economic risk and earning regional profits."
        )
    },

    "FRANCE_SUB": {
        "description": (
            "France is a limited-risk service provider responsible for sales and marketing support. "
            "It assists Luxembourg (and occasionally the US) with customer outreach and acquisition."
        ),
        "authority_limits": (
            "France has no authority to sign contracts on behalf of the group. "
            "This structure is intentionally designed to prevent the creation of a Permanent Establishment (PE) in France."
        ),
        "financial_mechanics": (
            "France earns a cost-plus service fee from Luxembourg and generates no independent revenue."
        ),
        "strategic_role": (
            "France functions strictly as a support entity rather than a revenue-earning business unit."
        )
    },

    "EU_CUSTOMERS": {
        "description": (
            "All EU customers contract directly with the Luxembourg entity. "
            "This centralizes commercial activity, aligns with the principal structure, "
            "and avoids accidental PE exposure in France."
        ),
        "revenue_flow": (
            "All EU revenue is recognized in Luxembourg, not in France."
        )
    },

    "MONEY_FLOWS": {
        "Luxembourg_to_US": (
            "Luxembourg pays royalties to the US for IP usage and sends dividends to the US once profitable."
        ),
        "Korea_to_US": (
            "Korea sends dividends upstream when profitable and provides R&D services to the US on a cost-plus basis."
        ),
        "Korea_to_Lux_or_US": (
            "Korea provides liquidity funding (loans or equity) to support group deficits."
        ),
        "France_to_Lux": (
            "France receives a cost-plus service fee from Luxembourg for sales support."
        )
    },

    "TAX_LOGIC": {
        "Transfer_Pricing": (
            "US owns the IP and receives royalties; Korea performs R&D and receives a cost-plus return; "
            "France provides sales support and receives a cost-plus fee; Luxembourg acts as the commercial principal "
            "and earns EU revenue."
        ),
        "PE_Risk": (
            "To avoid PE risk, France cannot sign contracts and does not negotiate binding terms. "
            "All contracts are executed in Luxembourg or the US."
        ),
        "Liquidity_Optimization": (
            "Korea holds most of the group’s cash and funds deficits in the US and Luxembourg. "
            "Future profits flow back to the US through royalties and dividends."
        ),
        "Future_Profit_Model": (
            "Luxembourg is strategically designed to become the EU profit center."
        )
    }
}


# -------------------------------
# KB File Helpers
# -------------------------------
def ensure_kb_exists():
    print("➡️ ensure_kb_exists() called")
    if not os.path.exists(KB_PATH):
        print("📁 KB file not found — creating new KB at:", KB_PATH)
        os.makedirs(os.path.dirname(KB_PATH), exist_ok=True)
        with open(KB_PATH, "w", encoding="utf-8") as f:
            json.dump(DEFAULT_KB, f, indent=2)
    else:
        print("✔ KB already exists")


def load_kb() -> Dict[str, Any]:
    print("➡️ load_kb() called")
    ensure_kb_exists()
    print("📄 Loading KB from:", KB_PATH)
    with open(KB_PATH, "r", encoding="utf-8") as f:
        return json.load(f)


def save_kb(kb: Dict[str, Any]) -> None:
    print("➡️ save_kb() called — writing KB to disk")
    os.makedirs(os.path.dirname(KB_PATH), exist_ok=True)
    with open(KB_PATH, "w", encoding="utf-8") as f:
        json.dump(kb, f, indent=2, ensure_ascii=False)


# -------------------------------
# Tool: fetch_kb()
# -------------------------------
@function_tool
def fetch_kb(ctx: RunContextWrapper[Any], section: str = None, query: str = None) -> str:
    """
    Returns the **entire knowledge base as a JSON string** for every agent call.
    Ignores section or query inputs.
    """
    print("\n==============================")
    print("📚 fetch_kb() TOOL CALLED — returning FULL KB")
    print("==============================")

    kb = load_kb()

    # Always return the full KB as JSON string
    return json.dumps(kb, indent=2, ensure_ascii=False)