"""
このモジュールは、画像ファイルをスキャンしテキスト化するためのユーティリティ関数です。

機能:
- ファイルをBase64形式にエンコードする関数
- 画像ファイルをテキスト化するための関数群
- テキスト化されたファイルをAPIに送信する関数
- ディレクトリ内のすべてのファイルまたは単一ファイルを処理するメイン関数
- CTRL+Cで現在処理中のテキスト化処理をキャンセルするリクエストを送信

使用例:
    python script_scanVisualDocuments.py <api_base_url> <auth_token> <directory_or_file_path> --model <model> --output_json

モジュール変数:
    SUCCESS_CODE (int): 成功時のステータスコード
    FAILURE_CODE (int): 失敗時のステータスコード
    URI_SCAN (str): 画像のテキスト化APIのエンドポイント
    URI_PROGRESS (str): テキスト化進捗確認APIのエンドポイント
    URI_RESULT (str): テキスト化結果取得APIのエンドポイント
    POLLING_INTERVAL (int): ポーリング時の待機時間（秒）
    FILE_INTERVAL (int): 複数ファイルを処理する際のファイル間の待機時間（秒）
    CERT (str or bool): SSL証明書の検証設定(SSL証明書の検証を行う場合はTrue、無効にする場合はFalseまたは証明書ファイルのパスを指定)

関数:
    encode_file_to_base64(file_path)
        ファイルをBase64エンコードします。

    scan(api_base_url, auth_token, file_path, model)
        画像ファイルをテキスト化処理のためにAPIに送信します。

    get_progress(api_base_url, request_id, auth_token)
        テキスト化処理の進捗状況を確認します。

    get_result(api_base_url, request_id, auth_token)
        テキスト化処理の結果を取得します。

    write_file(result, filename, output_json)
        テキスト化結果をMarkdownファイル, jsonファイルとして保存します。

    cancel_scan(api_base_url, request_id, auth_token)
        現在処理中のテキスト化処理をキャンセルします。

    scan_file(file_path, api_base_url, auth_token, model, output_json)
        画像ファイルをスキャンしてテキスト形式に変換します。

    check_extension(file_path)
        指定されたファイルパスの拡張子が対象の拡張子かチェックします。

    main(directory_or_file_path, api_base_url, auth_token, model, output_json)
        ディレクトリ内のファイルまたは単一ファイルに対して、スキャンとインデックス登録の全処理を実行します。
"""

import os
import sys
import base64
import requests
import json
import traceback
import urllib3
from urllib3.exceptions import InsecureRequestWarning
import time
from datetime import datetime
from zoneinfo import ZoneInfo
import argparse


# 成功時と失敗時(デフォルト)の返すコードをグローバル変数として定義
SUCCESS_CODE = 0
FAILURE_CODE = 1
CANCEL_CODE = 2
URI_SCAN = "/genai-api/v1/visualDocuments/scanAsync"
URI_PROGRESS = "/genai-api/v1/visualDocuments/scanStatus"
URI_RESULT = "/genai-api/v1/visualDocuments/scanResults"
URI_CANCEL = "/genai-api/v1/visualDocuments/scanCancel"
POLLING_INTERVAL = 10
FILE_INTERVAL = 10
CERT = True


def encode_file_to_base64(file_path: str) -> str:
    """
    ファイルをBase64エンコードします。

    Args:
        file_path (str): エンコードするファイルのパス。

    Returns:
        str: Base64エンコードされたファイルの文字列。
    """
    with open(file_path, "rb") as file:
        encoded_string = base64.b64encode(file.read()).decode("ascii")
    return encoded_string


def scan(
    api_base_url: str,
    auth_token: str,
    file_path: str,
    model: str,
) -> int:
    """
    画像ファイルをテキスト化処理のためにAPIに送信します。

    Args:
        api_base_url (str): APIのベースURL。
        auth_token (str): 認証トークン。
        file_path (str): 画像ファイルのパス。
        model (str): モデル名。

    Returns:
        tuple: (ステータスコード, リクエストID)
               成功時は(0, request_id)、失敗時は(エラーステータスコード, None)を返します。
    """
    encoded_file = encode_file_to_base64(file_path)
    filename = os.path.basename(file_path)

    payload = {
        "file": encoded_file,
        "filename": filename,
        "model": model,
    }

    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {auth_token}",
    }

    try:
        response = requests.post(api_base_url + URI_SCAN, headers=headers, json=payload, verify=CERT)
        now = datetime.now(ZoneInfo("Asia/Tokyo"))

        print(f"Processed {filename} at {now}: {response.status_code} {response.text}")

        if response.status_code != 202:
            try:
                data = response.json()
                print("Response Data:", data)
            except json.JSONDecodeError:
                print("Response is not in JSON format")
                print("Response text:", response.text)
            return response.status_code, None  # エラーステータスコードを返す

        data = response.json()
        print(f"request id: {data.get('id')}")
        return SUCCESS_CODE, data.get("id")  # 成功時

    except Exception as e:
        traceback.print_exc(file=sys.stdout)
        now = datetime.now(ZoneInfo("Asia/Tokyo"))
        print(f"[scan] Error scan document: {e}")
        print(f"Error occurred at: {now}")
        return FAILURE_CODE, None  # 失敗時のデフォルトエラーステータスコード


def get_progress(api_base_url: str, request_id: str, auth_token: str):
    """
    テキスト化処理の進捗状況を確認します。

    Args:
        api_base_url (str): APIのベースURL。
        request_id (str): テキスト化リクエストのID。
        auth_token (str): 認証トークン。

    Returns:
        tuple: (ステータスコード, 進捗データ)
               成功時は(0, データ)、失敗時は(エラーステータスコード, None)を返します。
    """
    headers = {
        "Authorization": f"Bearer {auth_token}",
    }

    try:
        response = requests.get(f"{api_base_url}{URI_PROGRESS}/{request_id}", headers=headers, verify=CERT)

        if response.status_code != 200:
            try:
                data = response.json()
                print("Response Data:", data)
            except json.JSONDecodeError:
                print("Response is not in JSON format")
                print("Response text:", response.text)
            return response.status_code, None  # エラーステータスコードを返す

        data = response.json()
        return SUCCESS_CODE, data  # 成功時

    except Exception as e:
        traceback.print_exc(file=sys.stdout)
        now = datetime.now(ZoneInfo("Asia/Tokyo"))
        print(f"[get_progress] Error scan document: {e}")
        print(f"Error occurred at: {now}")
        return FAILURE_CODE, None  # 失敗時のデフォルトエラーステータスコード


def get_result(api_base_url: str, request_id: str, auth_token: str):
    """
    テキスト化処理の結果を取得します。

    Args:
        api_base_url (str): APIのベースURL。
        request_id (str): テキスト化リクエストのID。
        auth_token (str): 認証トークン。

    Returns:
        tuple: (ステータスコード, 結果データ)
               成功時は(0, データ)、失敗時は(エラーステータスコード, None)を返します。
    """
    headers = {
        "Authorization": f"Bearer {auth_token}",
    }

    try:
        response = requests.get(f"{api_base_url}{URI_RESULT}/{request_id}", headers=headers, verify=CERT)

        if response.status_code != 200:
            try:
                data = response.json()
                print("Response Data:", data)
            except json.JSONDecodeError:
                print("Response is not in JSON format")
                print("Response text:", response.text)
            return response.status_code, None  # エラーステータスコードを返す

        data = response.json()
        return SUCCESS_CODE, data  # 成功時

    except Exception as e:
        traceback.print_exc(file=sys.stdout)
        now = datetime.now(ZoneInfo("Asia/Tokyo"))
        print(f"[get_result] Error scan document: {e}")
        print(f"Error occurred at: {now}")
        return FAILURE_CODE, None  # 失敗時のデフォルトエラーステータスコード


def write_file(result, filename="output", output_json=False):
    """
    テキスト化結果をMarkdownファイルとして保存します。

    Args:
        result (dict): テキスト化APIから返された結果データ。
        filename (str, optional): 出力ファイルのベース名。デフォルトは"output"。
        output_json (bool, optional): すべての結果をJSONでも出力するか。デフォルトはFalse。

    Returns:
        str: 保存されたファイルのパス。
    """
    pages = result["pages"]
    output_directory = "./output"
    
    # ファイル名から拡張子を除去し、.md拡張子を追加
    output_filename = f"{os.path.splitext(filename)[0]}.md"
    output_path = f"{output_directory}/{output_filename}"
    output_filename_json = f"{os.path.splitext(filename)[0]}.json"
    output_path_json = f"{output_directory}/{output_filename_json}"

    # ディレクトリが存在しない場合、作成する
    if not os.path.exists(output_directory):
        os.makedirs(output_directory)

    with open(output_path, "w", encoding="utf-8") as f:
        for page in pages:
            f.write("\n\n".join(chunk["text"] for chunk in page["chunks"]))
    print(f"Markdown file written to {output_path}")

    # JSONファイルの出力
    if output_json:
        # すべての結果を出力する
        with open(output_path_json, "w", encoding="utf-8") as f:
            json.dump(result, f, ensure_ascii=False, indent=2)
        print(f"All results written to {output_path_json}")

    return output_path


def cancel_scan(api_base_url: str, request_id: str, auth_token: str):
    """
    現在処理中のテキスト化処理をキャンセルします。

    Returns:
        tuple: ステータスコード
               成功時は0、失敗時はエラーステータスコードを返します。
    """
    payload = {
        "id": request_id,
    }

    headers = {
        "Authorization": f"Bearer {auth_token}",
    }

    try:
        response = requests.post(f"{api_base_url}{URI_CANCEL}", headers=headers, json=payload, verify=CERT)

        if response.status_code not in (200, 202):
            try:
                data = response.json()
                print("Cancel Response Data:", data)
            except json.JSONDecodeError:
                print("Cancel response is not in JSON format")
                print("Cancel response text:", response.text)
            return response.status_code

        return SUCCESS_CODE

    except Exception as e:
        traceback.print_exc(file=sys.stdout)
        now = datetime.now(ZoneInfo("Asia/Tokyo"))
        print(f"[cancel_scan] Error cancel scan: {e}")
        print(f"Error occurred at: {now}")
        return FAILURE_CODE


def scan_file(
    file_path: str,
    api_base_url: str,
    auth_token: str,
    model: str,
    output_json: bool,
) -> int:
    """
    画像ファイルをスキャンしてテキスト形式に変換します。

    Args:
        file_path (str): 画像ファイルのパス。
        api_base_url (str): APIのベースURL。
        auth_token (str): 認証トークン。
        model (str): モデル。
        output_json (bool): 結果をJSONでも出力するかどうか。

    Returns:
        tuple: (ステータスコード, 処理済みファイルパス)
               成功時は(0, 処理済みファイルパス, None)、失敗時は(エラーステータスコード, None, requestId)、キャンセル時は(2, None, requestId)を返します。
    """
    # 画像処理
    status_code, request_id = scan(api_base_url, auth_token, file_path, model)
    cancel_requested = False

    if status_code != SUCCESS_CODE:
        print("Failed to scan")
        return status_code, None, None

    # ステータス取得
    while True:
        try:
            status_code, data = get_progress(api_base_url, request_id, auth_token)
            if status_code != SUCCESS_CODE:
                print("Failed to get progress")
                return status_code, None, request_id

            status = data.get("status", "")
            cancelRequested = data.get("cancelRequested", False)
            print(f"  status: {status}, cancelRequested: {cancelRequested}")
            if status == "completed":
                completed_at = data.get("timestamp", 0)
                completed_at = datetime.fromtimestamp(completed_at, ZoneInfo("Asia/Tokyo"))
                print(f"Scanning completed: {file_path} at {completed_at}")
                break
            if status == "canceled":
                canceled_at = data.get("timestamp", 0)
                canceled_at = datetime.fromtimestamp(canceled_at, ZoneInfo("Asia/Tokyo"))
                print(f"Scanning canceled: {file_path} at {canceled_at}")
                return CANCEL_CODE, None, request_id
            if status == "failed":
                failed_at = data.get("timestamp", 0)
                failed_at = datetime.fromtimestamp(failed_at, ZoneInfo("Asia/Tokyo"))
                print(f"Scanning failed: {file_path} at {failed_at}")
                return FAILURE_CODE, None, request_id
            
            # 一定時間待機
            time.sleep(POLLING_INTERVAL)

        except KeyboardInterrupt:
            if not cancel_requested:
                cancel_requested = True
                print("Ctrl+C detected. Sending cancel request...")
                c_status = cancel_scan(api_base_url, request_id, auth_token)
                if c_status != SUCCESS_CODE:
                    print(f"Failed to send cancel request (status={c_status}).")
                    return CANCEL_CODE, None, request_id
                print("Cancel requested, please wait for the cancellation to complete.")
            continue

    # 結果取得
    status, result = get_result(api_base_url, request_id, auth_token)
    if status != SUCCESS_CODE:
        print("Failed to get result")
        return status, None, request_id
        
    # ファイル名を渡して出力ファイルを作成
    try:
        filename = os.path.basename(file_path)
        output_path = write_file(result, filename, output_json)
    except Exception as e:
        traceback.print_exc(file=sys.stdout)
        now = datetime.now(ZoneInfo("Asia/Tokyo"))
        print(f"[write_file] Error write file: {e}")
        print(f"Error occurred at: {now}")
        return FAILURE_CODE, None, request_id  # 失敗時のデフォルトエラーステータスコード

    return SUCCESS_CODE, output_path, None  # 全て成功時


def check_extension(file_path):
    """
    指定されたファイルパスの拡張子が対象の拡張子かチェックする
    対象: .pdf, .png, .jpg, .jpeg

    Args:
        file_path (str): チェックするファイルパス。
    
    Returns:
        bool: 対象の拡張子の場合は True、それ以外の場合は False。
    """
    allowed_extensions = {".pdf", ".png", ".jpg", ".jpeg"}
    extension = os.path.splitext(file_path)[1]
    return extension.lower() in allowed_extensions


def main(
    directory_or_file_path: str,
    api_base_url: str,
    auth_token: str,
    model: str,
    output_json: bool,
) -> int:
    """
    ディレクトリ内のファイルまたは単一ファイルに対して、スキャンとインデックス登録の全処理を実行します。

    Args:
        directory_or_file_path (str): ディレクトリまたはファイルのパス。
        api_base_url (str): APIのベースURL。
        auth_token (str): 認証トークン。
        model (str): モデル名。
        output_json (bool): 結果をJSONでも出力するかどうか。

    Returns:
        int: 成功時は0、失敗時はエラーステータスコードまたは1を返します。
    """
    # SSL証明書の検証を無効にする場合の警告を無効化
    if isinstance(CERT, bool) and not CERT:
        urllib3.disable_warnings(InsecureRequestWarning)  

    print(f"Starting process for directory or file: {directory_or_file_path}")
    print(f"Press Ctrl+C to cancel the current request.")

    if os.path.isdir(directory_or_file_path):
        scan_failed_list =[]
        failed_status_list = []
        canceled_status_list = []
        is_first = True
        for root, _, files in os.walk(directory_or_file_path):
            for file in files:
                file_path = os.path.join(root, file)
                
                # 拡張子チェック
                if not check_extension(file_path):
                    print(f"This file is skipped due to an unsupported extension: {file_path}")
                    continue

                if is_first:
                    is_first = False
                else:
                    # 負荷が上がりすぎないよう待機する
                    print(f"The processing of the next file will begin in {FILE_INTERVAL} seconds. Please standby.")
                    time.sleep(FILE_INTERVAL)

                print(f"Scanning file: {file_path}")
                status_code, _, request_id = scan_file(file_path, api_base_url, auth_token, model, output_json)
                if status_code == CANCEL_CODE:
                    print(f"Scan canceled: {file_path}")
                    canceled_status_list.append((request_id, file_path))
                    continue
                elif status_code != SUCCESS_CODE:
                    print(f"Failed to scan: {file_path}")
                    if request_id:
                        failed_status_list.append((request_id, file_path))
                    else:
                        scan_failed_list.append((status_code, file_path))
                    continue

        print("Processing completed.")

        # 失敗結果出力
        if scan_failed_list:
            print("========== List of Failed Scan Requests ==========")
            for scan_failed in scan_failed_list:
                print(f"  FilePath: {scan_failed[1]}, StatusCode: {scan_failed[0]} ")
        # リトライ処理ありの失敗
        if failed_status_list:
            print("========== List of Failed Scan Processes ==========")
            for failed_status in failed_status_list:
                print(f"  FilePath: {failed_status[1]}, RequestID: {failed_status[0]} ")
        # キャンセル
        if canceled_status_list:
            print("========== List of Canceled Scan Processes ==========")
            for canceled_status in canceled_status_list:
                print(f"  FilePath: {canceled_status[1]}, RequestID: {canceled_status[0]} ")

    elif os.path.isfile(directory_or_file_path):
        # 単一ファイルの場合は、直接そのファイルパスを使用
        file_path = directory_or_file_path

        # 拡張子チェック
        if not check_extension(file_path):
            print(f"Unsupported file extension: {file_path}")
            return FAILURE_CODE

        print(f"Scanning file: {file_path}")
        status_code, _, request_id = scan_file(file_path, api_base_url, auth_token, model, output_json)

        if status_code == CANCEL_CODE:
            print(f"Scan canceled: {file_path}")
            return CANCEL_CODE
        elif status_code != SUCCESS_CODE:
            if request_id:
                print(f"  FilePath: {file_path}, RequestID: {request_id} ")
            else:
                print(f"  FilePath: {file_path}, StatusCode: {status_code} ")
            return status_code
        
    else:
        print(f"Error: {directory_or_file_path} is neither a file nor a directory")
        return FAILURE_CODE

    return SUCCESS_CODE  # 全て成功時


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Process some files.')
    parser.add_argument('api_base_url', type=str, help='API Base URL')
    parser.add_argument('auth_token', type=str, help='Authentication token')
    parser.add_argument('directory_or_file_path', type=str, help='Path to the directory or file')
    parser.add_argument('--model', type=str, choices=["scan-std-model-v1-jp", "scan-std-model-v1-apac", "scan-std-model-v2-apac"], default="scan-std-model-v2-apac", help='Model name')
    parser.add_argument('--output_json', default=False, help='Output all results', action="store_true")

    args = parser.parse_args()

    exit_code = main(
        args.directory_or_file_path,
        args.api_base_url,
        args.auth_token,
        model=args.model,
        output_json=args.output_json,
    )

    sys.exit(exit_code)