# 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_create_classifier.py

DESCRIPTION:
    This sample demonstrates how to create a classifier analyzer to categorize documents and use it
    to analyze documents with and without automatic segmentation.

    ## About classifiers

    Classifiers are a type of custom analyzer that create classification workflows to categorize
    documents into predefined custom categories using ContentCategories. They allow you to perform
    classification and content extraction as part of a single API call. Classifiers are useful for:
    - Content organization: Organize large document collections by type through categorization
    - Data routing (optional): Optionally route your data to specific custom analyzers based on
      category, ensuring your data is routed to the best analyzer for processing when needed
    - Multi-document processing: Process files containing multiple document types by automatically
      segmenting them

    Classifiers use custom categories to define the types of documents they can identify. Each
    category has a Description that helps the AI model understand what documents belong to that
    category. You can define up to 200 category names and descriptions. You can include an "other"
    category to handle unmatched content; otherwise, all files are forced to be classified into one
    of your defined categories.

    The enable_segment property in the analyzer configuration controls whether multi-document files
    are split into segments:
    - enable_segment = False: Classifies the entire file as a single category (classify only)
    - enable_segment = True: Automatically splits the file into segments by category (classify and
      segment)

USAGE:
    python sample_create_classifier.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).

    Before using classifiers, you MUST configure model deployments for your Microsoft Foundry
    resource. See sample_update_defaults.py for setup instructions.
"""

import os
import time
from typing import cast

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

load_dotenv()


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

    client = ContentUnderstandingClient(endpoint=endpoint, credential=credential)

    # [START create_classifier]
    # Generate a unique analyzer ID
    analyzer_id = f"my_classifier_{int(time.time())}"

    print(f"Creating classifier '{analyzer_id}'...")

    # Define content categories for classification.
    # Each category has a description that helps the AI model identify matching documents.
    # Optionally, set analyzer_id on a category to route matched segments to a prebuilt
    # or custom analyzer for field extraction. For example, setting
    # analyzer_id="prebuilt-invoice" on the Invoice category will automatically extract
    # invoice fields (vendor, line items, totals, etc.) from segments classified as Invoice.
    categories = {
        "Loan_Application": ContentCategoryDefinition(
            description="Documents submitted by individuals or businesses to request funding, "
            "typically including personal or business details, financial history, "
            "loan amount, purpose, and supporting documentation."
        ),
        "Invoice": ContentCategoryDefinition(
            description="Billing documents issued by sellers or service providers to request "
            "payment for goods or services, detailing items, prices, taxes, totals, "
            "and payment terms.",
            analyzer_id="prebuilt-invoice",  # Route Invoice segments for field extraction
        ),
        "Bank_Statement": ContentCategoryDefinition(
            description="Official statements issued by banks that summarize account activity "
            "over a period, including deposits, withdrawals, fees, and balances."
        ),
    }

    # Create analyzer configuration
    config = ContentAnalyzerConfig(
        return_details=True,
        enable_segment=True,  # Enable automatic segmentation by category
        content_categories=categories,
    )

    # Create the classifier analyzer
    classifier = ContentAnalyzer(
        base_analyzer_id="prebuilt-document",
        description="Custom classifier for financial document categorization",
        config=config,
        models={"completion": "gpt-4.1"},
    )

    # Create the classifier
    poller = client.begin_create_analyzer(
        analyzer_id=analyzer_id,
        resource=classifier,
    )
    result = poller.result()  # Wait for creation to complete

    # Get the full analyzer details after creation
    result = client.get_analyzer(analyzer_id=analyzer_id)

    print(f"Classifier '{analyzer_id}' created successfully!")
    if result.description:
        print(f"  Description: {result.description}")
    # [END create_classifier]

    # [START analyze_with_classifier]
    file_path = "sample_files/mixed_financial_docs.pdf"

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

    print(f"\nAnalyzing document with classifier '{analyzer_id}'...")

    analyze_poller = client.begin_analyze_binary(
        analyzer_id=analyzer_id,
        binary_input=file_bytes,
    )
    analyze_result: AnalysisResult = analyze_poller.result()

    # Display classification results
    if analyze_result.contents and len(analyze_result.contents) > 0:
        document_content = cast(DocumentContent, analyze_result.contents[0])
        print(
            f"Pages: {document_content.start_page_number}-{document_content.end_page_number}"
        )

        # Display segments (classification results)
        if document_content.segments and len(document_content.segments) > 0:
            print(f"\nFound {len(document_content.segments)} segment(s):")
            for segment in document_content.segments:
                print(f"  Category: {segment.category or '(unknown)'}")
                print(f"  Pages: {segment.start_page_number}-{segment.end_page_number}")
                print(f"  Segment ID: {segment.segment_id or '(not available)'}")
                print()
        else:
            print("No segments found (document classified as a single unit).")
    else:
        print("No content found in the analysis result.")
    # [END analyze_with_classifier]

    # [START classifier_to_llm_input]
    # Convert classification results to LLM-friendly text.
    # to_llm_input automatically detects classification results: it expands the parent
    # into per-segment blocks, each with its category label in the YAML front matter.
    # Segments are separated by a ***** divider.
    print("\n" + "=" * 60)
    print("CLASSIFICATION RESULT AS LLM INPUT")
    print("=" * 60)

    text = to_llm_input(analyze_result)
    print(text)
    # [END classifier_to_llm_input]

    # Clean up - delete the classifier
    print(f"\nCleaning up: deleting classifier '{analyzer_id}'...")
    client.delete_analyzer(analyzer_id=analyzer_id)
    print(f"Classifier '{analyzer_id}' deleted successfully.")


if __name__ == "__main__":
    main()
