# Pytanie: "Jaka jest stolica Polski?"
# Agent może:
# 1. Odpowiedzieć z własnej wiedzy (bez Bing) → brak citations
# 2. Użyć Bing Search → z citations

message.content[0].text.annotations[0].url_citation.url
message.content[0].text.annotations[0].url_citation.title

run = project_client.agents.runs.create_and_process(
    thread_id=thread.id,
    agent_id=agent.id,
    tool_choice={"type": "bing_grounding"}  # Wymuś Bing!
)

#!/usr/bin/env python3
"""
Tool choice i Citations extraction z Bing Search
"""

from azure.ai.projects import AIProjectClient
from azure.identity import ChainedTokenCredential, EnvironmentCredential, AzureCliCredential
import json

print("✓ Importy załadowane")

# Konfiguracja (zastąp swoimi wartościami!)
PROJECT_ENDPOINT = "YOUR_PROJECT_ENDPOINT_HERE"  # Z Lab 07
AGENT_ID = "YOUR_AGENT_WITH_BING_ID_HERE"  # Agent z Bing tool

# Utworzenie klienta
print("\n🔧 Tworzenie klienta...")
credential = ChainedTokenCredential(
    EnvironmentCredential(),
    AzureCliCredential()
)

# (Opcjonalne) Jeżeli masz rozszyfrowywanie ruchu na proxy (SSL interception),
# możesz wyłączyć weryfikowanie certyfikatu:
# from azure.core.pipeline.transport import RequestsTransport
# transport = RequestsTransport(connection_verify=False)

project_client = AIProjectClient(
    endpoint=PROJECT_ENDPOINT,
    credential=credential
    # transport=transport  # Odkomentuj jeżeli używasz proxy
)

# Utworzenie thread
print("🔧 Tworzenie thread...")
thread = project_client.agents.threads.create()

print("✓ Klient i thread gotowe")
print(f"  Thread ID: {thread.id}")

# Test 1: Auto mode - agent sam decyduje
print("\n" + "=" * 60)
print("TEST 1: AUTO MODE (agent decyduje)")
print("=" * 60)

# Pytanie które agent MOŻE odpowiedzieć bez Bing
question1 = "Jaka jest stolica Polski?"

# Dodaj wiadomość
message1 = project_client.agents.messages.create(
    thread_id=thread.id,
    role="user",
    content=question1
)

# Run BEZ tool_choice (agent sam decyduje)
print(f"\n▶️  Pytanie: {question1}")
print("   Tool choice: AUTO (agent decyduje)")
run1 = project_client.agents.runs.create_and_process(
    thread_id=thread.id,
    agent_id=AGENT_ID
    # Brak tool_choice = auto mode
)

# Odczyt odpowiedzi
messages1 = project_client.agents.messages.list(thread_id=thread.id)
for msg in messages1:
    if msg.role == "assistant" and msg.run_id == run1.id:
        response1 = msg.content[0].text.value
        break

print("\n✓ Odpowiedź agenta:")
print("-" * 60)
print(response1[:300])  # Pierwsze 300 znaków
print("-" * 60)

# Sprawdź czy użyto Bing
run_steps1 = project_client.agents.run_steps.list(
    thread_id=thread.id,
    run_id=run1.id
)

used_bing = False
for step in run_steps1:
    if hasattr(step, 'step_details') and hasattr(step.step_details, 'tool_calls'):
        for tool_call in step.step_details.tool_calls:
            if tool_call.type == "bing_grounding":
                used_bing = True
                break

print(f"\n📊 Użyto Bing Search: {'✓ TAK' if used_bing else '✗ NIE (z własnej wiedzy)'}")
print(f"📊 Tokeny: {run1.usage.total_tokens if run1.usage else 'N/A'}")

# Test 2: Forced Bing - wymuszenie użycia Bing
print("\n" + "=" * 60)
print("TEST 2: FORCED BING (tool_choice wymusza)")
print("=" * 60)

# TO SAMO pytanie!
question2 = "Jaka jest stolica Polski?"

# Dodaj wiadomość
message2 = project_client.agents.messages.create(
    thread_id=thread.id,
    role="user",
    content=question2
)

# Run Z tool_choice - wymuszamy Bing!
print(f"\n▶️  Pytanie: {question2}")
print("   Tool choice: FORCED Bing")
run2 = project_client.agents.runs.create_and_process(
    thread_id=thread.id,
    agent_id=AGENT_ID,
    tool_choice={"type": "bing_grounding"}  # Wymuś Bing!
)

# Odczyt odpowiedzi
response2 = None
messages2 = project_client.agents.messages.list(thread_id=thread.id)
for msg in messages2:
    if msg.role == "assistant" and msg.run_id == run2.id:
        response2 = msg.content[0].text.value
        break

if response2 is None:
    response2 = "Brak odpowiedzi od agenta."
    print("Brak odpowiedzi od agenta.")

print("\n✓ Odpowiedź agenta:")
print("-" * 60)
print(response2[:400])  # Pierwsze 400 znaków
print("-" * 60)

# Sprawdź czy użyto Bing (powinno być TAK!)
run_steps2 = project_client.agents.run_steps.list(
    thread_id=thread.id,
    run_id=run2.id
)

used_bing2 = False
for step in run_steps2:
    if hasattr(step, 'step_details') and hasattr(step.step_details, 'tool_calls'):
        for tool_call in step.step_details.tool_calls:
            if tool_call.type == "bing_grounding":
                used_bing2 = True
                break

print(f"\n📊 Użyto Bing Search: {'✓ TAK (wymuszono!)' if used_bing2 else '✗ NIE'}")
print(f"📊 Tokeny: {run2.usage.total_tokens if run2.usage else 'N/A'}")

# Ekstrakcja citations z Test 2
print("\n" + "=" * 60)
print("EKSTRAKCJA CITATIONS")
print("=" * 60)

# Pobierz message z Test 2
messages_with_citations = project_client.agents.messages.list(thread_id=thread.id)
for msg in messages_with_citations:
    if msg.role == "assistant" and msg.run_id == run2.id:
        # Sprawdź czy są annotations (citations)
        text_content = msg.content[0]

        if hasattr(text_content, 'text') and hasattr(text_content.text, 'annotations'):
            annotations = text_content.text.annotations

            if annotations and len(annotations) > 0:
                print(f"\n✓ Znaleziono {len(annotations)} citations:")
                print("-" * 60)

                for i, annotation in enumerate(annotations, 1):
                    # Bing citations są typu url_citation (NIE uri_citation!)
                    if hasattr(annotation, 'url_citation'):
                        citation = annotation.url_citation
                        print(f"\n[{i}] {citation.title}")
                        print(f"    URL: {citation.url}")

                        # Text który został zacytowany
                        if hasattr(annotation, 'text'):
                            print(f"    Ref: {annotation.text}")
            else:
                print("\n✗ Brak citations (agent nie użył Bing)")

        break

print("-" * 60)

# Test 3: Citations w JSON format (schema bez format enforcement)
print("\n" + "=" * 60)
print("TEST 3: CITATIONS W JSON (schema only)")
print("=" * 60)

# Pytanie wymagające Bing
question3 = "Jaka jest dzisiejsza pogoda w Warszawie?"

# Dodaj wiadomość
message3 = project_client.agents.messages.create(
    thread_id=thread.id,
    role="user",
    content=question3
)

# Run z Bing + JSON format (schema bez "no markdown")
print(f"\n▶️  Pytanie: {question3}")
print("   Tool choice: FORCED Bing")
print("   Response format: JSON (schema only)")
run3 = project_client.agents.runs.create_and_process(
    thread_id=thread.id,
    agent_id=AGENT_ID,
    tool_choice={"type": "bing_grounding"},
    response_format={"type": "json_object"},
    additional_instructions="""Return ONLY raw JSON object with this exact schema:
{
  "location": "city name",
  "weather": "current weather description",
  "sources": ["list of source URLs"]
}
No explanations, no markdown, no text before JSON. Pure JSON only."""
)

# Odczyt odpowiedzi
response3 = None
messages3 = project_client.agents.messages.list(thread_id=thread.id)
for msg in messages3:
    if msg.role == "assistant" and msg.run_id == run3.id:
        response3 = msg.content[0].text.value
        break
if response3 is None:
    response3 = "Brak odpowiedzi od agenta."
    print("Brak odpowiedzi od agenta.")

print("\n✓ Odpowiedź agenta (JSON):")
print("-" * 60)
print(response3[:400])
print("-" * 60)

# Próba parsowania
print("\n🔍 Próba parsowania JSON:")
try:
    # Strip markdown if present
    json_text = response3
    if response3.strip().startswith("```"):
        import re
        json_match = re.search(r'```json\s*(\{.*\})\s*```', response3, re.DOTALL)
        if json_match:
            json_text = json_match.group(1)

    data3 = json.loads(json_text)
    print("✓ Parse SUCCESS")
    print(json.dumps(data3, indent=2, ensure_ascii=False)[:300])
except json.JSONDecodeError as e:
    print(f"✗ Parse FAILED: {e}")

print(f"\n📊 Tokeny: {run3.usage.total_tokens if run3.usage else 'N/A'}")

def extract_citations(message):
    """Extract citations from agent message"""
    citations = []

    if hasattr(message.content[0], 'text') and hasattr(message.content[0].text, 'annotations'):
        for annotation in message.content[0].text.annotations:
            if hasattr(annotation, 'url_citation'):
                citations.append({
                    'title': annotation.url_citation.title,
                    'url': annotation.url_citation.url,
                    'text': annotation.text if hasattr(annotation, 'text') else ''
                })

    return citations

# Usage
citations = extract_citations(message3)
for i, cite in enumerate(citations, 1):
    print(f"[{i}] {cite['title']}")
    print(f"    {cite['url']}")

tool_choice={"type": "bing_grounding"}  # Correct
# NIE: tool_choice="bing"
# NIE: tool_choice={"bing_grounding": True}

# W Azure AI Foundry Portal:
# Agents → Twój agent → Tools → sprawdź "Bing Grounding"

# BAD: "Wymyśl nową bajkę" - nie da się wyszukać!
# GOOD: "Jaka jest pogoda w Warszawie?" - da się wyszukać

annotations = message.content[0].text.annotations
if annotations:
    for ann in annotations:
        if hasattr(ann, 'url_citation'):
            print(ann.url_citation.title)
else:
    print("No annotations found")

# Markdown = native citations z title + URL
run = project_client.agents.runs.create_and_process(
    thread_id=thread.id,
    agent_id=AGENT_ID,
    tool_choice={"type": "bing_grounding"}
    # BEZ response_format!
)
# Potem extract z annotations

# Auto (default) - agent decyduje
run = project_client.agents.runs.create_and_process(
    thread_id=thread.id,
    agent_id=agent.id
)

# Forced - MUSISZ użyć Bing
run = project_client.agents.runs.create_and_process(
    thread_id=thread.id,
    agent_id=agent.id,
    tool_choice={"type": "bing_grounding"}
)

# None - NIE MOŻESZ użyć narzędzi
run = project_client.agents.runs.create_and_process(
    thread_id=thread.id,
    agent_id=agent.id,
    tool_choice={"type": "none"}
)

message.content[0].text.annotations[i].url_citation.title  # "Wikipedia - Warszawa"
message.content[0].text.annotations[i].url_citation.url    # "https://pl.wikipedia.org/..."
message.content[0].text.annotations[i].text                # "[1]"