import io, base64, math
import uuid

from PIL import Image, ImageOps, ImageFilter
from openai import AsyncOpenAI
from pathlib import Path
from pdf2image import convert_from_path
import asyncio

# 请求模型地址
model_base_url = ""

async def img_to_base64_qwen3vl(original, path, max_side=2096):
    img = Image.open(path)
    img = ImageOps.exif_transpose(img).convert("RGB")

    w, h = img.size
    s = min(1.0, max_side / max(w, h))
    if s < 1:
        img = img.resize((int(w * s), int(h * s)), Image.Resampling.LANCZOS)

    img = img.filter(ImageFilter.UnsharpMask(1, 120, 3))

    w, h = img.size
    nw, nh = math.ceil(w / 32) * 32, math.ceil(h / 32) * 32
    img = ImageOps.expand(
        img,
        ((nw - w) // 2, (nh - h) // 2, nw - w - (nw - w) // 2, nh - h - (nh - h) // 2),
        fill=(128, 128, 128)
    )

    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=90, subsampling=0, optimize=True)
    b64 = base64.b64encode(buf.getvalue()).decode()
    img.save(f"./SC/{str(uuid.uuid4())}_{original}", format="JPEG", quality=90, subsampling=0, optimize=True)
    return f"data:image/jpeg;base64,{b64}"


def pdf_to_images_sync(original, pdf_path, dpi=300, fmt="png"):
    try:
        poppler_path = r"C:\Users\FLYF\Downloads\Release-25.12.0-0\poppler-25.12.0\Library\bin"

        pdf_path = Path(pdf_path)
        id = str(uuid.uuid4())

        output_dir = Path(f"./TEMP/{pdf_path.stem}_{id}"[:200])
        output_dir.mkdir(parents=True, exist_ok=True)

        images = convert_from_path(
            pdf_path,
            dpi=dpi,
            poppler_path=poppler_path
        )

        img_paths = []
        ima_names = []

        for i, img in enumerate(images, start=1):
            out_path = output_dir / f"{pdf_path.stem}_{i}.{fmt}"
            img.save(out_path)
            img_paths.append(str(out_path))
            ima_names.append(f"{original}_{i}.{fmt}")

        # output = {
        #     "pdf_name": original,
        #     "result": img_paths,
        #     "state": "success"
        # }

        return img_paths, ima_names
    except Exception as e:
        return img_paths, ima_names


async def pdf_to_images_async(*args, **kwargs):
    return await asyncio.to_thread(pdf_to_images_sync, *args, **kwargs)


with open("SystemPrompt", 'r', encoding='utf-8') as f:
    SystemPrompt = f.read()

client = AsyncOpenAI(api_key="1", base_url=model_base_url)


async def qwen3vl(original, img_path, prompt):
    try:
        img = await img_to_base64_qwen3vl(original, img_path)
        response = await client.chat.completions.create(
            model="Qwen3-VL-FLYFAI",
            messages=[
                {"role": "system", "content": SystemPrompt},
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": img
                            }
                        },
                        {
                            "type": "text",
                            "text": prompt
                        }
                    ]
                }
            ],
            temperature=0.1,
            top_p=0.3,
            max_tokens=4096,
            timeout=90
        )
        output = {
            "image_name": original,
            "result": response.choices[0].message.content,
            "state": "success"
        }
        return output
    except Exception as e:
        return {"image_name": original, "result": str(e), "state": "error"}


async def run_one(semaphore, prompt, original, img_path, retries=3):
    async with semaphore:
        for i in range(retries):
            try:
                return await asyncio.wait_for(qwen3vl(original, img_path, prompt), timeout=90)
            except Exception as e:
                if i == retries - 1:
                    return {"image_name": original, "result": str(e), "state": "error"}
                await asyncio.sleep(3)


async def run_images(original, img_paths, prompt, max_concurrency=30):
    sem = asyncio.Semaphore(max_concurrency)
    tasks = [run_one(sem, prompt, original, p) for original, p in zip(original, img_paths)]
    return await asyncio.gather(*tasks)


async def run_pdf(original, pdf_path, prompt):
    img_paths, ima_names = await pdf_to_images_async(original, pdf_path)
    result = await run_images(ima_names, img_paths, prompt)
    return result