# pylint: disable=line-too-long,useless-suppression
# coding=utf-8
# --------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for license information.
# --------------------------------------------------------------------------
"""
FILE: sample_analyze_binary_async.py

DESCRIPTION:
    This sample demonstrates how to analyze a PDF file from disk using the prebuilt-documentSearch
    analyzer and convert the result to LLM-friendly text.

    ## About analyzing documents from binary data

    One of the key values of Content Understanding is taking a content file and extracting the content
    for you in one call. The service returns an AnalysisResult that contains an array of AnalysisContent
    items in AnalysisResult.contents. This sample starts with a document file, so each item is a
    DocumentContent (a subtype of AnalysisContent) that exposes markdown plus detailed structure such
    as pages, tables, figures, and paragraphs.

    ## Using results with LLMs

    The markdown returned by Content Understanding can be directly consumed by large language models
    (LLMs) for summarization, question answering, and other generative AI tasks. To make this even
    easier, the SDK provides a convenient ``to_llm_input()`` helper that converts an AnalysisResult
    into a single text block with YAML front matter (content type, page numbers, extracted fields)
    followed by the markdown body — ready for injection into LLM prompts, vector databases, or
    agentic tool outputs. For advanced usage (output options, content ranges, video/audio, metadata),
    see sample_to_llm_input_async.py.

    This sample focuses on document analysis. For prebuilt RAG analyzers covering images, audio, and
    video, see sample_analyze_url_async.py.

    ## Prebuilt analyzers

    Content Understanding provides prebuilt RAG analyzers (the prebuilt-*Search analyzers, such as
    prebuilt-documentSearch) that return markdown and a one-paragraph Summary for each content item,
    making them useful for retrieval-augmented generation (RAG) and other downstream applications:

    - prebuilt-documentSearch - Extracts content from documents (PDF, images, Office documents) with
      layout preservation, table detection, figure analysis, and structured markdown output.
      Optimized for RAG scenarios.
    - prebuilt-audioSearch - Transcribes audio content with speaker diarization, timing information,
      and conversation summaries. Supports multilingual transcription.
    - prebuilt-videoSearch - Analyzes video content with visual frame extraction, audio transcription,
      and structured summaries. Provides temporal alignment of visual and audio content.
    - prebuilt-imageSearch - Analyzes standalone images and returns a one-paragraph Summary of the
      image content. For images that contain text (including hand-written text), use
      prebuilt-documentSearch.

    This sample uses prebuilt-documentSearch to extract structured content from PDF documents.

USAGE:
    python sample_analyze_binary_async.py

    Set the environment variables with your own values before running the sample:
    1) CONTENTUNDERSTANDING_ENDPOINT - the endpoint to your Content Understanding resource.
    2) CONTENTUNDERSTANDING_KEY - your Content Understanding API key (optional if using DefaultAzureCredential).

    See sample_update_defaults_async.py for model deployment setup guidance.
"""

import asyncio
import os

from dotenv import load_dotenv
from azure.ai.contentunderstanding import to_llm_input
from azure.ai.contentunderstanding.aio import ContentUnderstandingClient
from azure.ai.contentunderstanding.models import (
    AnalysisResult,
    DocumentContent,
)
from azure.core.credentials import AzureKeyCredential
from azure.identity.aio import DefaultAzureCredential

load_dotenv()


async def main() -> None:
    endpoint = os.environ["CONTENTUNDERSTANDING_ENDPOINT"]
    key = os.getenv("CONTENTUNDERSTANDING_KEY")
    credential = AzureKeyCredential(key) if key else DefaultAzureCredential()

    async with ContentUnderstandingClient(
        endpoint=endpoint, credential=credential
    ) as client:
        # [START analyze_document_from_binary]
        # Replace with the path to your local document file.
        file_path = "sample_files/sample_invoice.pdf"

        with open(file_path, "rb") as f:
            file_bytes = f.read()

        print(f"Analyzing {file_path} with prebuilt-documentSearch...")
        poller = await client.begin_analyze_binary(
            analyzer_id="prebuilt-documentSearch",
            binary_input=file_bytes,
        )
        result: AnalysisResult = await poller.result()
        # [END analyze_document_from_binary]

        # [START analyze_binary_with_content_range]
        # Use a multi-page document for content range demonstrations.
        multi_page_path = "sample_files/mixed_financial_invoices.pdf"
        with open(multi_page_path, "rb") as f:
            multi_page_bytes = f.read()

        # Analyze only pages 3 onward.
        print("\nAnalyzing pages 3 onward with content range '3-'...")
        range_poller = await client.begin_analyze_binary(
            analyzer_id="prebuilt-documentSearch",
            binary_input=multi_page_bytes,
            content_range="3-",
        )
        range_result: AnalysisResult = await range_poller.result()

        if isinstance(range_result.contents[0], DocumentContent):
            range_doc = range_result.contents[0]
            page_nums = [p.page_number for p in range_doc.pages] if range_doc.pages else []
            print(f"Content range analysis returned pages {page_nums}")
        # [END analyze_binary_with_content_range]

        # [START analyze_binary_with_combined_content_range]
        # Analyze pages 1-3, page 5, and pages 9 onward.
        print("\nAnalyzing combined pages (1-3, 5, 9-) with content range '1-3,5,9-'...")
        combine_range_poller = await client.begin_analyze_binary(
            analyzer_id="prebuilt-documentSearch",
            binary_input=multi_page_bytes,
            content_range="1-3,5,9-",
        )
        combine_range_result: AnalysisResult = await combine_range_poller.result()

        if isinstance(combine_range_result.contents[0], DocumentContent):
            combine_doc = combine_range_result.contents[0]
            page_nums = [p.page_number for p in combine_doc.pages] if combine_doc.pages else []
            print(f"Combined content range analysis returned pages {page_nums}")
        # [END analyze_binary_with_combined_content_range]

        # [START convert_to_llm_input]
        # The markdown above can be consumed directly by LLMs. For convenience, the SDK
        # provides to_llm_input() which packages the result into a single text block with
        # YAML front matter (content type, pages, fields, optional metadata) followed by
        # the markdown body — ready for LLM prompts, vector stores, or agentic tools.
        print("\n" + "=" * 60)
        print("LLM-READY OUTPUT")
        print("=" * 60)

        text = to_llm_input(result)
        print(text)
        # [END convert_to_llm_input]

        # [START extract_markdown]
        print("\n" + "=" * 60)
        print("MARKDOWN CONTENT")
        print("=" * 60)

        # A PDF file has only one content element even if it contains multiple pages
        content = result.contents[0]
        print(content.markdown)
        # [END extract_markdown]

        # [START access_document_properties]
        print("\n" + "=" * 60)
        print("DOCUMENT PROPERTIES")
        print("=" * 60)

        # Check if this is document content to access document-specific properties
        if isinstance(content, DocumentContent):
            print(f"\nDocument type: {content.mime_type or '(unknown)'}")
            print(f"Start page: {content.start_page_number}")
            print(f"End page: {content.end_page_number}")

            # Check for pages
            if content.pages and len(content.pages) > 0:
                print(f"\nNumber of pages: {len(content.pages)}")
                for page in content.pages:
                    unit = content.unit or "units"
                    print(
                        f"  Page {page.page_number}: {page.width} x {page.height} {unit}"
                    )

            # Check for tables
            if content.tables and len(content.tables) > 0:
                print(f"\nNumber of tables: {len(content.tables)}")
                table_counter = 1
                for table in content.tables:
                    print(
                        f"  Table {table_counter}: {table.row_count} rows x {table.column_count} columns"
                    )
                    table_counter += 1
        # [END access_document_properties]

    if not isinstance(credential, AzureKeyCredential):
        await credential.close()


if __name__ == "__main__":
    asyncio.run(main())
