%%capture
# Foundry IQ rides on the public-preview azure-search-documents SDK, which
# exposes the 2026-05-01-preview Knowledge Base / Knowledge Source surface.
# It ships on PyPI -- no extra package feed required.
import importlib.metadata as _md

try:
    _ok = _md.version("azure-search-documents").startswith("12.1.0b")
except _md.PackageNotFoundError:
    _ok = False

if not _ok:
    %pip install --quiet \
        "azure-search-documents==12.1.0b1" \
        "azure-identity>=1.19.0" \
        "azure-core>=1.32.0" \
        "python-dotenv>=1.0.1"

import os
from typing import Optional

from azure.core.credentials import AzureKeyCredential
from azure.search.documents import __version__ as search_sdk_version
from azure.search.documents.indexes import SearchIndexClient
from dotenv import load_dotenv

load_dotenv(override=True)


def env(name: str, *, required: bool = True, default: Optional[str] = None) -> Optional[str]:
    value = os.getenv(name, default)
    if required and not value:
        raise RuntimeError(f"Missing required env var: {name}")
    return value or None


def foundry_resource_uri(endpoint: str) -> str:
    # Strip any /openai/... path so we have the bare Foundry resource URI.
    return endpoint.split("/openai/", 1)[0].rstrip("/")


# ---- Required ------------------------------------------------------------
SEARCH_ENDPOINT = env("SEARCH_ENDPOINT")
SEARCH_API_KEY = env("SEARCH_API_KEY")
AOAI_ENDPOINT = foundry_resource_uri(env("AOAI_ENDPOINT"))
AOAI_API_KEY = env("AOAI_API_KEY")

GPT_DEPLOYMENT = env("AOAI_GPT_DEPLOYMENT", required=False, default="gpt-4.1")
GPT_MODEL = env("AOAI_GPT_MODEL", required=False, default=GPT_DEPLOYMENT)

# ---- Resource names ------------------------------------------------------
KS_WORK_IQ = "miq-ks-work-iq"
KS_FABRIC_IQ = "miq-ks-fabric-iq"
KS_WEB_IQ = "miq-ks-web-iq"
KB_NAME = "miq-foundry-iq-kb"

credential = AzureKeyCredential(SEARCH_API_KEY)
index_client = SearchIndexClient(endpoint=SEARCH_ENDPOINT, credential=credential)

# ---- Trackers consumed by the KB section and cleanup --------------------
created_ks: list[str] = []
created_resources: dict[str, str] = {}

print(f"azure-search-documents : {search_sdk_version}")
print(f"Search service         : {SEARCH_ENDPOINT}")
print(f"Foundry endpoint       : {AOAI_ENDPOINT}")
print(f"Chat deployment        : {GPT_DEPLOYMENT} ({GPT_MODEL})")

from azure.core.exceptions import ResourceNotFoundError


def skip(section: str, reason: str) -> None:
    print(f"[skipped] {section}: {reason}")


def summarize_ks(name: str) -> None:
    """Print the SDK status of a Knowledge Source after create."""
    try:
        status = index_client.get_knowledge_source_status(name)
    except Exception as exc:  # pragma: no cover - best-effort summary
        print(f"  status: <unavailable: {exc.__class__.__name__}>")
        return
    sync = getattr(status, "synchronization_status", None) or "?"
    last_state = getattr(status, "last_synchronization_state", None)
    last = getattr(last_state, "status", None) if last_state else "n/a"
    stats = getattr(status, "statistics", None) or {}
    print(f"  synchronizationStatus={sync}  lastSync={last}  stats={stats}")

from azure.search.documents.indexes.models import WorkIQKnowledgeSource

# Work IQ takes no parameters object -- just name + description (kind="workIQ").
ks = WorkIQKnowledgeSource(
    name=KS_WORK_IQ,
    description="Work IQ -- tenant mail, chats, files, and meetings via Microsoft 365.",
)
try:
    index_client.create_or_update_knowledge_source(ks)
    created_ks.append(KS_WORK_IQ)
    print(f"Created {KS_WORK_IQ}")
    summarize_ks(KS_WORK_IQ)
except Exception as exc:
    skip(KS_WORK_IQ, f"create failed (tenant may not be entitled for Work IQ): {exc}")

from azure.search.documents.indexes.models import (
    FabricOntologyKnowledgeSource,
    FabricOntologyKnowledgeSourceParameters,
)

# Point this at YOUR ontology. Both IDs are in the Fabric ontology item URL:
#   https://app.fabric.microsoft.com/groups/<workspace-id>/ontologies/<ontology-id>
# This recipe uses an airline ontology as the running example; substitute your own.
FABRIC_WORKSPACE_ID = env("FABRIC_WORKSPACE_ID", required=False)
FABRIC_ONTOLOGY_ID = env("FABRIC_ONTOLOGY_ID", required=False)

if not (FABRIC_WORKSPACE_ID and FABRIC_ONTOLOGY_ID):
    skip(KS_FABRIC_IQ, "set FABRIC_WORKSPACE_ID and FABRIC_ONTOLOGY_ID to your own ontology to enable Fabric IQ")
else:
    ks = FabricOntologyKnowledgeSource(
        name=KS_FABRIC_IQ,
        description="Fabric IQ -- a business ontology (here, an airline ontology: flights, routes, aircraft, crews) in Microsoft Fabric.",
        fabric_ontology_parameters=FabricOntologyKnowledgeSourceParameters(
            workspace_id=FABRIC_WORKSPACE_ID,
            ontology_id=FABRIC_ONTOLOGY_ID,
        ),
    )
    try:
        index_client.create_or_update_knowledge_source(ks)
        created_ks.append(KS_FABRIC_IQ)
        print(f"Created {KS_FABRIC_IQ} (workspace={FABRIC_WORKSPACE_ID}, ontology={FABRIC_ONTOLOGY_ID})")
        summarize_ks(KS_FABRIC_IQ)
    except Exception as exc:
        skip(KS_FABRIC_IQ, f"create failed (check Search MI access to the Fabric workspace): {exc}")

from azure.search.documents.indexes.models import (
    McpServerKnowledgeSource,
    McpServerKnowledgeSourceParameters,
    McpServerStoredHeadersAuthentication,
    McpServerStoredHeadersParameters,
)

WEB_IQ_MCP_SERVER_URL = "https://api.microsoft.ai/v3/mcp"  # public URL -- fine to commit

# The API key is a SECRET: read it from the environment, never hard-code it.
web_iq_key = env("WEB_IQ_MCP_API_KEY", required=False)

if not web_iq_key:
    skip(KS_WEB_IQ, "set WEB_IQ_MCP_API_KEY (your Web IQ MCP key) to enable Web IQ -- join the waitlist at https://aka.ms/webiq-waitlist")
else:
    web_iq_tools = [
        {"name": "web",    "outputParsing": {"kind": "auto"}, "inclusionMode": "reranked", "maxOutputTokens": 4096},
        {"name": "news",   "outputParsing": {"kind": "auto"}, "inclusionMode": "reranked", "maxOutputTokens": 4096},
        {"name": "videos", "outputParsing": {"kind": "auto"}, "inclusionMode": "reranked", "maxOutputTokens": 4096},
        {"name": "browse", "outputParsing": {"kind": "auto"}, "inclusionMode": "reranked", "maxOutputTokens": 4096},
    ]
    ks = McpServerKnowledgeSource(
        name=KS_WEB_IQ,
        description="Web IQ -- Microsoft Grounding MCP server (web, news, videos, browse).",
        mcp_server_parameters=McpServerKnowledgeSourceParameters(
            server_url=WEB_IQ_MCP_SERVER_URL,
            # storedHeaders auth: send the Web IQ x-apikey directly.
            authentication=McpServerStoredHeadersAuthentication(
                stored_headers_parameters=McpServerStoredHeadersParameters(
                    headers={"x-apikey": web_iq_key},
                ),
            ),
            tools=web_iq_tools,
        ),
    )
    try:
        index_client.create_or_update_knowledge_source(ks)
        created_ks.append(KS_WEB_IQ)
        print(f"Created {KS_WEB_IQ} ({WEB_IQ_MCP_SERVER_URL})")
        summarize_ks(KS_WEB_IQ)
    except Exception as exc:
        skip(KS_WEB_IQ, f"create failed: {exc}")

from azure.search.documents.indexes.models import (
    AzureOpenAIVectorizerParameters,
    KnowledgeBase,
    KnowledgeBaseAzureOpenAIModel,
    KnowledgeSourceReference,
)
from azure.search.documents.knowledgebases.models import (
    KnowledgeRetrievalLowReasoningEffort,
    KnowledgeRetrievalOutputMode,
)

if not created_ks:
    raise RuntimeError(
        "No Knowledge Sources were created -- cannot build a KB. Enable at least "
        "one of Work IQ / Fabric IQ / Web IQ above."
    )

# Build the Foundry IQ KB over whatever Microsoft IQ sources came up this run.
kb_sources = list(created_ks)
print(f"Building Foundry IQ KB '{KB_NAME}' over {len(kb_sources)} federated source(s):")
for n in kb_sources:
    print(f"  - {n}")

gpt_params = AzureOpenAIVectorizerParameters(
    resource_url=AOAI_ENDPOINT,
    deployment_name=GPT_DEPLOYMENT,
    api_key=AOAI_API_KEY,
    model_name=GPT_MODEL,
)

kb = KnowledgeBase(
    name=KB_NAME,
    description="Foundry IQ KB federating Work IQ, Fabric IQ (airline ontology), and Web IQ.",
    models=[KnowledgeBaseAzureOpenAIModel(azure_open_ai_parameters=gpt_params)],
    knowledge_sources=[KnowledgeSourceReference(name=n) for n in kb_sources],
    retrieval_reasoning_effort=KnowledgeRetrievalLowReasoningEffort(),
    output_mode=KnowledgeRetrievalOutputMode.ANSWER_SYNTHESIS,
    retrieval_instructions=(
        "Route each subquery to the Microsoft IQ most likely to answer it: "
        "use Work IQ for internal, people, and collaboration context; use "
        "Fabric IQ for business facts from the ontology (fleet, routes, "
        "operations); use Web IQ for current events and public information. "
        "Use several sources when a question spans more than one."
    ),
    answer_instructions=(
        "Answer using only the retrieved content across all sources. "
        "When a question spans work content, the airline ontology, and the web, "
        "integrate them faithfully and preserve the [ref_id:N] citations."
    ),
)
index_client.create_or_update_knowledge_base(kb)
created_resources["kb"] = KB_NAME
print(f"\nFoundry IQ Knowledge Base '{KB_NAME}' is ready.")

from azure.identity import DefaultAzureCredential
from azure.search.documents.knowledgebases import KnowledgeBaseRetrievalClient
from azure.search.documents.knowledgebases.models import (
    KnowledgeBaseMessage,
    KnowledgeBaseMessageTextContent,
    KnowledgeBaseRetrievalRequest,
)

kb_client = KnowledgeBaseRetrievalClient(
    endpoint=SEARCH_ENDPOINT,
    knowledge_base_name=KB_NAME,
    credential=DefaultAzureCredential(),
)

request = KnowledgeBaseRetrievalRequest(
    messages=[
        KnowledgeBaseMessage(
            role="user",
            content=[KnowledgeBaseMessageTextContent(
                text="What should I know before our transatlantic route-planning review?"
            )],
        ),
    ],
)

# Work IQ + Fabric IQ are per-user; pass the caller's token. Web IQ uses its own key.
user_token = DefaultAzureCredential().get_token("https://search.azure.com/.default").token
result = kb_client.retrieve(request, query_source_authorization=user_token)

print(result.response[0].content[0].text)

# 1) Delete the KB.
if created_resources.get("kb"):
    try:
        index_client.delete_knowledge_base(created_resources["kb"])
        print(f"Deleted KB {created_resources['kb']}")
    except ResourceNotFoundError:
        pass
    except Exception as exc:
        print(f"KB delete: {exc}")

# 2) Delete every KS we created (Work IQ, Fabric IQ, Web IQ).
for ks_name in created_ks:
    try:
        index_client.delete_knowledge_source(ks_name)
        print(f"Deleted KS {ks_name}")
    except ResourceNotFoundError:
        pass
    except Exception as exc:
        print(f"KS {ks_name} delete: {exc}")