BIPAuto/utils/discrete_material_plan/extractor.py

"""
Discrete Material Plan Data Extractor

Pure functions for extracting and post-processing discrete material plan data.
All functions are stateless and accept required parameters explicitly.
Caller is responsible for browser/session lifecycle management.
"""

import pandas as pd
from pathlib import Path
from typing import List, Optional, Tuple
from playwright.sync_api import Page, Frame


def chunk_order_ids(order_ids: List[str], batch_size: int) -> List[List[str]]:
    """
    Split order IDs into batches.

    Args:
        order_ids: List of order IDs to process
        batch_size: Maximum number of order IDs per batch

    Returns:
        List of batches, where each batch is a list of order IDs

    Example:
        >>> chunk_order_ids(["A", "B", "C", "D"], 2)
        [["A", "B"], ["C", "D"]]
    """
    return [
        order_ids[i : i + batch_size]
        for i in range(0, len(order_ids), batch_size)
    ]


def get_login_url(base_url: str) -> str:
    """
    Construct the complete login URL from base URL.

    Args:
        base_url: Base ERP URL (e.g., "https://erp.example.com")

    Returns:
        Complete login page URL

    Example:
        >>> get_login_url("https://erp.example.com")
        "https://erp.example.com/yonbip/resources/uap/rbac/login/main/index.html"
    """
    return f"{base_url.rstrip('/')}/yonbip/resources/uap/rbac/login/main/index.html"


def extract_batch(
    work_frame: Frame,
    page: Page,
    order_ids: List[str],
    batch_index: int,
    download_dir: str,
) -> str:
    """
    Execute download workflow for a single batch of order IDs.

    This is a thin wrapper around `execute_batch_download_workflow` from
    extractor_core.py, providing progress reporting context.

    Args:
        work_frame: The work iframe containing the search form and data grid
        page: The Playwright page object for download handling
        order_ids: List of order IDs for this batch
        batch_index: Zero-based batch index for naming the output file
        download_dir: Directory path to save the downloaded file

    Returns:
        Full path to the downloaded Excel file
    """
    from .extractor_core import execute_batch_download_workflow

    return execute_batch_download_workflow(
        work_frame=work_frame,
        page=page,
        order_ids=order_ids,
        batch_index=batch_index,
        download_dir=download_dir,
    )


def extract_batches(
    work_frame: Frame,
    page: Page,
    order_ids: List[str],
    download_dir: str,
    batch_size: int = 10,
) -> List[str]:
    """
    Download data for multiple batches of order IDs.

    Caller is responsible for:
    - Browser session management (login, logout)
    - Navigation to discrete material plan page
    - Query interface setup

    Args:
        work_frame: The work iframe containing the data grid
        page: The Playwright page object for download handling
        order_ids: List of order IDs to download
        download_dir: Directory path to save downloaded files
        batch_size: Maximum number of order IDs per batch

    Returns:
        List of paths to downloaded Excel files

    Example:
        >>> # Caller manages session
        >>> browser, context, page, main_frame = login(...)
        >>> work_frame, page1 = navigate_to_discrete_material_page(main_frame, page)
        >>> setup_query_interface(work_frame)
        >>> files = extract_batches(work_frame, page1, order_ids, "/downloads")
        >>> context.close()
        >>> browser.close()
    """
    from .extractor_core import setup_query_interface

    downloaded_files = []
    chunks = chunk_order_ids(order_ids, batch_size)

    # Setup query interface once
    setup_query_interface(work_frame)

    # Process each batch
    for batch_index, batch in enumerate(chunks):
        file_path = extract_batch(
            work_frame=work_frame,
            page=page,
            order_ids=batch,
            batch_index=batch_index,
            download_dir=download_dir,
        )
        downloaded_files.append(file_path)

    return downloaded_files


def post_process_downloads(
    downloaded_files: List[str],
    output_file: str,
    verbose: bool = True,
    cleanup_temp_files: bool = True,
) -> Tuple[str, pd.DataFrame]:
    """
    Convert and merge downloaded Excel files into structured DataFrame.

    Uses ExcelConverter to convert each file, then merges all results.

    Args:
        downloaded_files: List of paths to downloaded Excel files
        output_file: Path to save merged Excel result
        verbose: Whether to print progress messages
        cleanup_temp_files: Whether to delete temporary downloaded files after processing (default: True)

    Returns:
        Tuple of (output_file_path, merged_dataframe)

    Example:
        >>> output_path, df = post_process_downloads(
        ...     downloaded_files=["batch_1.xlsx", "batch_2.xlsx"],
        ...     output_file="merged.xlsx",
        ...     verbose=True,
        ...     cleanup_temp_files=True
        ... )
    """
    from .excel_converter import ExcelConverter

    converter = ExcelConverter(verbose=verbose)
    all_dfs = []

    # Convert each file
    for i, file_path in enumerate(downloaded_files):
        if verbose:
            print(f"Converting file {i + 1}/{len(downloaded_files)}: {file_path}")

        # Convert (do not save intermediate result)
        df = converter.convert(input_file=file_path)
        all_dfs.append(df)

    # Merge all DataFrames
    if not all_dfs:
        merged_df = pd.DataFrame()
    else:
        merged_df = pd.concat(all_dfs, ignore_index=True)

    # Save merged result
    output_path = Path(output_file)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    merged_df.to_excel(output_path, index=False)

    if verbose:
        print(f"Merged result saved to: {output_path}")
        print(f"Total rows: {len(merged_df)}")

    # Cleanup temporary downloaded files
    if cleanup_temp_files:
        _cleanup_temp_files(downloaded_files, verbose)

    return str(output_path), merged_df


def _cleanup_temp_files(downloaded_files: List[str], verbose: bool = True) -> int:
    """
    Remove temporary downloaded files.

    Args:
        downloaded_files: List of file paths to delete
        verbose: Whether to print progress messages

    Returns:
        Number of files successfully deleted
    """
    deleted_count = 0
    for file_path in downloaded_files:
        try:
            Path(file_path).unlink()
            deleted_count += 1
            if verbose:
                print(f"Deleted temp file: {file_path}")
        except Exception as e:
            if verbose:
                print(f"Warning: Could not delete {file_path}: {e}")
    return deleted_count


def extract_and_post_process(
    work_frame: Frame,
    page: Page,
    order_ids: List[str],
    download_dir: str,
    output_file: str,
    batch_size: int = 10,
    verbose: bool = True,
    cleanup_temp_files: bool = True,
) -> Tuple[str, pd.DataFrame]:
    """
    Complete extraction workflow: download batches + post-process to merged Excel.

    This is a high-level convenience function that orchestrates the full workflow.
    Caller is still responsible for browser session management.

    Args:
        work_frame: The work iframe containing the data grid
        page: The Playwright page object for download handling
        order_ids: List of order IDs to extract
        download_dir: Directory for temporary batch files
        output_file: Path for final merged Excel output
        batch_size: Maximum order IDs per batch
        verbose: Whether to print progress messages
        cleanup_temp_files: Whether to delete temporary downloaded files after processing (default: True)

    Returns:
        Tuple of (output_file_path, merged_dataframe)

    Example:
        >>> # Caller manages session
        >>> browser, context, page, main_frame = login(...)
        >>> work_frame, page1 = navigate_to_discrete_material_page(main_frame, page)
        >>> output_path, df = extract_and_post_process(
        ...     work_frame, page1, order_ids, "/downloads", "output.xlsx",
        ...     cleanup_temp_files=True
        ... )
        >>> context.close()
        >>> browser.close()
    """
    # Step 1: Download all batches
    if verbose:
        print(f"Downloading {len(order_ids)} orders in batches of {batch_size}...")

    downloaded_files = extract_batches(
        work_frame=work_frame,
        page=page,
        order_ids=order_ids,
        download_dir=download_dir,
        batch_size=batch_size,
    )

    if verbose:
        print(f"Downloaded {len(downloaded_files)} batch file(s)")

    # Step 2: Post-process (convert + merge)
    output_path, merged_df = post_process_downloads(
        downloaded_files=downloaded_files,
        output_file=output_file,
        verbose=verbose,
        cleanup_temp_files=cleanup_temp_files,
    )

    return output_path, merged_df


def read_order_ids_from_file(id_file: str, encoding: str = "utf-8") -> List[str]:
    """
    Read order IDs from a text file (one ID per line).

    Args:
        id_file: Path to file containing order IDs
        encoding: File encoding (default: utf-8)

    Returns:
        List of order IDs (stripped, empty lines filtered)

    Raises:
        FileNotFoundError: If id_file doesn't exist

    Example:
        >>> order_ids = read_order_ids_from_file("orders.txt")
    """
    id_path = Path(id_file)
    if not id_path.exists():
        raise FileNotFoundError(f"Order ID file not found: {id_file}")

    with open(id_path, "r", encoding=encoding) as f:
        return [line.strip() for line in f if line.strip()]


def extract_from_file(
    id_file: str,
    work_frame: Frame,
    page: Page,
    download_dir: str,
    output_file: str,
    batch_size: int = 10,
    verbose: bool = True,
    cleanup_temp_files: bool = True,
) -> Tuple[str, pd.DataFrame]:
    """
    Extract data from order IDs in a file and post-process to merged Excel.

    Convenience function that reads IDs from file and calls extract_and_post_process.

    Args:
        id_file: Path to file containing order IDs (one per line)
        work_frame: The work iframe containing the data grid
        page: The Playwright page object for download handling
        download_dir: Directory for temporary batch files
        output_file: Path for final merged Excel output
        batch_size: Maximum order IDs per batch
        verbose: Whether to print progress messages
        cleanup_temp_files: Whether to delete temporary downloaded files after processing (default: True)

    Returns:
        Tuple of (output_file_path, merged_dataframe)

    Example:
        >>> # Read order IDs
        >>> order_ids = read_order_ids_from_file("orders.txt")
        >>> # Extract and process
        >>> output_path, df = extract_from_file(
        ...     "orders.txt", work_frame, page, "/downloads", "output.xlsx",
        ...     cleanup_temp_files=True
        ... )
    """
    order_ids = read_order_ids_from_file(id_file)

    if verbose:
        print(f"Loaded {len(order_ids)} order IDs from {id_file}")

    return extract_and_post_process(
        work_frame=work_frame,
        page=page,
        order_ids=order_ids,
        download_dir=download_dir,
        output_file=output_file,
        batch_size=batch_size,
        verbose=verbose,
        cleanup_temp_files=cleanup_temp_files,
    )