import os
import sys
import json
try:
    from dotenv import load_dotenv
except:
    print(f"\nImport error: dotenv")
    input("Install: pip install dotenv\n")
    exit()
try:
    import openai
    from openai import OpenAI
except:
    print(f"\nImport error: openai")
    input("Install: pip install openai\n")
    exit()


config_path = "translate.env"
input_path = "transcript.txt"
output_path = ""
guess_speakers = 1

max_bytes = 20000
max_tokens = 2000

prompt = """
あなたは半導体・材料科学の教授です。
以下の講義の会話テキストには誤字・誤訳が多くあります。次の作業をしてください。
１．誤字・誤訳を訂正する
２．読みやすいように文ごとに改行を入れる
３．話者を推測する
"""


if not os.path.isfile(config_path):
    script_dir = os.path.dirname(os.path.abspath(__file__))
    config_path = os.path.join(script_dir, config_path)

print()
# 環境変数読み込み
if os.path.isfile(config_path):
    print(f"config_path: {config_path}")
else:
    print(f"Warning: config_path {config_path} is not found")
load_dotenv(dotenv_path=config_path)

account_inf_path = os.getenv("account_inf_path", "accounts.env")
if os.path.isfile(account_inf_path):
    print(f"account_inf_path: {account_inf_path}")
else:
    print(f"Warning: account_inf_path {account_inf_path} is not found")
load_dotenv(dotenv_path=account_inf_path)

# APIキー設定
api_key = os.getenv("OPENAI_API_KEY")
if api_key is None or api_key == "":
    print(f"Error: api_key is not found.")
    exit()

client = OpenAI(api_key=api_key)

openai_model = os.getenv("openai_model", "gpt-4o")
temperature = float(os.getenv("temperature", "0.3"))
max_tokens = int(os.getenv("max_tokens", max_tokens))


argv = sys.argv
nargs = len(argv)
if nargs > 1: input_path = argv[1]
if nargs > 2: output_path = argv[2]
if nargs > 3: max_bytes = int(argv[3])

# speaker guess用Function Call関数定義
functions = [
    {
        "name": "guess_speakers",
        "description": "会話のテキストから話者を推測し、JSON 構造で返す",
        "parameters": {
            "type": "object",
            "properties": {
                "speakers": {
                    "type": "array",
                    "items": {"type": "string"}
                },
                "turns": {
                    "type": "array",
                    "items": {
                        "type": "object",
                        "properties": {
                            "speaker": {"type": "string"},
                            "text": {"type": "string"}
                        },
                        "required": ["speaker", "text"]
                    }
                }
            },
            "required": ["speakers", "turns"]
        }
    }
]

def print_speakers(json_str: str):
    data = json.loads(json_str)
    for turn in data.get("turns", []):
        speaker = turn.get("speaker", "Unknown")
        text    = turn.get("text", "")
        print(f"[{speaker}]: {text}")

def text_with_speakers(json_str: str):
    data = json.loads(json_str)
    text_ret = ""
    for turn in data.get("turns", []):
        speaker = turn.get("speaker", "Unknown")
        text    = turn.get("text", "")
        text_ret += f"[{speaker}]: {text}\n"

    return text_ret

def save_with_speakers(json_str: str, outfile: str):
    """
    json_str に含まれる "turns" 配列を
      [SpeakerX]: 発言テキスト
    の形式で outfile に書き出します。
    """
    data = json.loads(json_str)
    with open(outfile, "w", encoding="utf-8") as f:
        for turn in data.get("turns", []):
            speaker = turn.get("speaker", "Unknown")
            text    = turn.get("text", "")
            f.write(f"[{speaker}]: {text}\n")

def save_text_to_file(filename: str, content: str):
    """
    指定された内容をファイルに保存
    """
    try:
        with open(filename, "w", encoding="utf-8") as f:
            f.write(content)
        print(f"'{filename}' に内容を保存しました。")
    except IOError as e:
        print(f"ファイルの保存中にエラーが発生しました: {e}")

def main():
    global output_path

    print()
    print("Guess speakers from text file")

    output_dir = os.path.dirname(os.path.abspath(input_path)) or "."
    file_body = os.path.splitext(os.path.basename(input_path))[0]

    #input_pathの存在を確認したのち、全テキストをtranscriptに読み込む
    if os.path.exists(input_path):
        with open(input_path, "r", encoding="utf-8") as f:
            transcript = f.read()
    else:
        print(f"\nError: Could not read [{input_path}]\n")
        exit()

    if output_path == "":
        output_path = os.path.join(output_dir, f"{file_body}_with_speakers.txt")    

    print()
    print(f"api_key: {api_key}")
    print(f"input_path  : {input_path}")
    print(f"output_path : {output_path}")
    print(f"guess_speakers: {guess_speakers}")
    print(f"openai_model: {openai_model}")
    print(f"temperature : {temperature}")
    print(f"max_tokens  : {max_tokens}")
    print(f"max_bytes   : {max_bytes}")

    if True:
        # GPT に Function Calling 形式でリクエスト
        print()
        print("Sending request to ChatGPT:")
        response = client.chat.completions.create(
            model="gpt-4o",
            messages=[
                {"role": "system", "content": prompt},
                {"role": "user",   "content": transcript[:max_bytes]}
            ],
            functions=functions,
            function_call={"name": "guess_speakers"},
            temperature=0.0
        )

        # 関数呼び出しの arguments 部分を JSON としてロード
        try:
            arguments = response.choices[0].message.function_call.arguments
            result = json.loads(arguments)
        except Exceptions as e:
            print(f"Error: JSONの解析中にエラーが起こりました: {e}")
#        print("0=", response.choices[0])
            print("response.choices[0].message: ", response.choices[0].message)
            print("response.choices[0].message.function_call:", response.choices[0].message.function_call)
            exit()
            
#        base = os.path.splitext(os.path.basename(audio_path))[0]
#        out_file = os.path.join(os.path.dirname(os.path.abspath(audio_path)), f"{base}_guessed.json")

        print_speakers(arguments)
        save_with_speakers(arguments, output_path)

        """
        # JSON 出力
        with open(out_file, "w", encoding="utf-8") as f:
            json.dump(result, f, ensure_ascii=False, indent=2)
        """
    
        print()
        print(f"話者推測結果を '{output_path}' に保存しました。")

    input("\nPress ENTER to terminate>>\n")
    

if __name__ == "__main__":
    main()
