import json
import os
import threading
import time

from pilot.job.impl.base_job import BaseJob

from pilot.generater.vera import VeraSingleton

class GenerateJsonBaseJob(BaseJob):

    prompt_content: str
    result_content: str
    result_file_path: str
    stream = False

    def run(self):
        prompt = self.prompt_content
        # トークン数チェック
        vera_ai = VeraSingleton.get_instance()
        #token_count = vertexai.count_tokens(prompt)
        #if token_count == 0:
        #    super().run()
        #    return
        #if token_count > 900000:
        #    print(f"警告: promptのトークン数が900000を超えています ({token_count} tokens)")
        #    super().run()
        #    return
        # VertexAI で生成
        start = time.time()
        result = vera_ai.generate_content(prompt,stream=self.stream)
        try:
            result_content = result.get('response', '')
            data = json.loads(result_content)
        except json.decoder.JSONDecodeError:
            result = vera_ai.generate_content(prompt)
            result_content = result.get('response', '')
            data = json.loads(result_content)
        end = time.time()
        print(f"Ai 処理時間 {self.file_path}: {end - start:.2f}秒")
        with open(self.result_file_path, 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
        super().run()