import os
import base64
import time
from dotenv import load_dotenv
from openai import OpenAI
import httpx
from PIL import Image

load_dotenv()

MODELS_TO_TEST = [
    "qwen3.5-122b"
]

client = OpenAI(
    api_key=os.getenv("API_KEY"),
    base_url=os.getenv("LAOZHANG_BASE_URL"),
    http_client=httpx.Client(timeout=httpx.Timeout(600.0, connect=30.0))
)


def encode_image(image_path):
    """Кодирует изображение в base64, конвертируя TIFF в JPEG при необходимости"""
    if image_path.lower().endswith('.tif') or image_path.lower().endswith('.tiff'):
        with Image.open(image_path) as img:
            if img.mode not in ('RGB', 'L'):
                img = img.convert('RGB')
            import io
            buffer = io.BytesIO()
            img.save(buffer, format='JPEG', quality=95)
            buffer.seek(0)
            return base64.b64encode(buffer.read()).decode('utf-8')
    else:
        with open(image_path, "rb") as img:
            return base64.b64encode(img.read()).decode('utf-8')


def get_instructions(criteria_file):
    with open("SKILL.md", "r", encoding="utf-8") as f:
        skill = f.read()
    c_path = os.path.join("references", criteria_file)
    with open(c_path, "r", encoding="utf-8") as f:
        criteria = f.read()
    return f"{skill}\n\n{criteria}"


def run_mass_check(base_dir="photo", criteria_file="russian-essay-criteria.md"):
    students = [d for d in os.listdir(
        base_dir) if os.path.isdir(os.path.join(base_dir, d))]

    if not students:
        print("[!] В папке photo пусто. Создай там папки с именами учеников.")
        return

    print(f"=== ЗАПУСК МАССОВОЙ ПРОВЕРКИ: {len(students)} учеников ===")
    instructions = get_instructions(criteria_file)

    for student in students:
        student_path = os.path.join(base_dir, student)
        photos = sorted([f for f in os.listdir(student_path)
                        if f.lower().endswith(('.jpg', '.jpeg', '.png', '.tif', '.tiff'))])

        if not photos:
            print(f"[SKIP] У {student} нет фото.")
            continue

        print(f"\n>>> РАБОТАЕМ С: {student.upper()} ({len(photos)} листа)")

        message_content = [
            {"type": "text", "text": "Распознай рукописный текст и проверь сочинение строго по критериям ФИПИ."}]

        for p in photos:
            b64 = encode_image(os.path.join(student_path, p))
            message_content.append({
                "type": "image_url",
                "image_url": {"url": f"data:image/jpeg;base64,{b64}"}
            })

        for model_id in MODELS_TO_TEST:
            safe_name = model_id.replace("/", "_")

            # ✅ ОТРЕДАКТИРОВАНО: отчёт сохраняется ВНУТРЬ папки ученика
            output_file = os.path.join(student_path, f"REPORT_{safe_name}.md")

            old_file = os.path.join(
                student_path, f"REPORT_{student}_{safe_name}.md")

            new_file = os.path.join(student_path, f"REPORT_{safe_name}.md")

            if os.path.exists(old_file) or os.path.exists(new_file):
                print(f"  [-] {model_id}: Уже проверено.")
                continue

            print(f"  [!] Запуск {model_id}...")
            start_time = time.time()

            try:
                response = client.chat.completions.create(
                    model=model_id,
                    messages=[
                        {"role": "system", "content": instructions},
                        {"role": "user", "content": message_content}
                    ],
                    temperature=0.0
                )

                res_text = response.choices[0].message.content
                duration = round(time.time() - start_time, 1)

                with open(output_file, "w", encoding="utf-8") as f:
                    header = f"--- \n**Ученик:** {student}\n**Модель:** {model_id}\n**Время:** {duration} сек.\n---\n\n"
                    f.write(header + res_text)

                print(f"  [OK] Готово! ({duration} сек.)")
                print(f"  📁 Сохранено: {output_file}")
                time.sleep(12)

            except Exception as e:
                print(f"  [ERR] Ошибка у {model_id}: {str(e)}")


if __name__ == "__main__":
    run_mass_check()
    print("\n=== ВСЕ ПРОВЕРКИ ВЫПОЛНЕНЫ ===")