import os
import base64
import asyncio
import time
import io
from dotenv import load_dotenv
from openai import AsyncOpenAI  # Переходим на АСИНХРОННЫЙ
import httpx
from PIL import Image

load_dotenv()

# Оставляем только Qwen, как ты и хотел
MODELS_TO_TEST = ["qwen3.5-122b"]

# Асинхронный клиент с запасом по таймауту
client = AsyncOpenAI(
    api_key=os.getenv("QWEN_API_KEY"),
    base_url=os.getenv("QWEN_BASE_URL"),
    http_client=httpx.AsyncClient(timeout=httpx.Timeout(900.0, connect=60.0))
)


def encode_image(image_path):
    """Сжимает и кодирует изображение, чтобы избежать Connection Error"""
    with Image.open(image_path) as img:
        # Конвертируем в RGB если нужно
        if img.mode not in ('RGB', 'L'):
            img = img.convert('RGB')

        # Уменьшаем размер до разумного (max 1600px по широкой стороне)
        # Это критически важно для стабильности соединения!
        img.thumbnail((1600, 1600))

        buffer = io.BytesIO()
        # Сжатие 85% - золотая середина
        img.save(buffer, format='JPEG', quality=85)
        return base64.b64encode(buffer.getvalue()).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}"


async def process_student(student, student_path, instructions):
    """Логика проверки одного студента"""
    # Поиск фото
    photos = sorted([f for f in os.listdir(student_path)
                    if f.lower().endswith(('.jpg', '.jpeg', '.png', '.tif', '.tiff'))])
    if not photos:
        return

    # ПОИСК ИСХОДНОГО ТЕКСТА (source.txt)
    source_text = ""
    source_path = os.path.join(student_path, "source.txt")
    if os.path.exists(source_path):
        with open(source_path, "r", encoding="utf-8") as f:
            source_text = f.read()

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

    prompt_text = "Распознай текст и проверь сочинение строго по критериям."
    if source_text:
        prompt_text += f"\n\nИСХОДНЫЙ ТЕКСТ ДЛЯ СВЕРКИ:\n{source_text}"

    message_content = [{"type": "text", "text": prompt_text}]

    # Кодируем фото (теперь со сжатием!)
    for p in photos:
        try:
            b64 = encode_image(os.path.join(student_path, p))
            message_content.append({
                "type": "image_url",
                "image_url": {"url": f"data:image/jpeg;base64,{b64}"}
            })
        except Exception as e:
            print(f"  [!] Ошибка обработки фото {p}: {e}")

    for model_id in MODELS_TO_TEST:
        safe_name = model_id.replace("/", "_")
        output_file = os.path.join(student_path, f"REPORT_{safe_name}.md")

        # 1. Формируем базовое имя, которое мы ХОТИМ создать
        safe_name = model_id.replace("/", "_")
        output_file = os.path.join(student_path, f"REPORT_{safe_name}.md")

        # 2. УМНЫЙ ПОИСК: проверяем, есть ли ВООБЩЕ любой файл, начинающийся на REPORT_ и содержащий имя модели
        existing_reports = [f for f in os.listdir(student_path)
                            if f.startswith("REPORT_") and safe_name in f]

        if existing_reports:
            print(
                f"  [-] {model_id}: Уже есть отчет ({existing_reports[0]}). Пропускаю.")
            continue

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

        try:
            # Асинхронный вызов
            response = await 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} сек.)")
            await asyncio.sleep(10)  # Пауза между запросами

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


async 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:
        return

    instructions = get_instructions(criteria_file)

    # Обрабатываем по одному, чтобы не перегружать канал и не ловить Connection Error
    for student in students:
        await process_student(student, os.path.join(base_dir, student), instructions)

if __name__ == "__main__":
    asyncio.run(run_mass_check())