import os
import time
import requests
from dotenv import load_dotenv

load_dotenv()

API_URL = os.getenv('OCR_API_URL')
HEADERS = {
    "Authorization": f"bearer {os.getenv('OCR_API_TOKEN')}",
    "Content-Type": "application/json",
}
MODEL = os.getenv("OCR_MODEL")


def ocr_from_url(image_url: str) -> str:
    r = requests.post(
        API_URL,
        json={
            "fileUrl": image_url,
            "model": MODEL,
            "optionalPayload": {
                "useDocOrientationClassify": False,
                "useDocUnwarping": False,
                "useChartRecognition": False,
            },
        },
        headers=HEADERS,
    )
    r.raise_for_status()
    job_id = r.json()["data"]["jobId"]

    while True:
        r = requests.get(f"{API_URL}/{job_id}", headers=HEADERS)
        r.raise_for_status()
        data = r.json()["data"]

        if data["state"] == "done":
            jsonl_url = data["resultUrl"]["jsonUrl"]
            break
        if data["state"] == "failed":
            raise RuntimeError("OCR failed")

        time.sleep(3)

    lines = requests.get(jsonl_url).text.strip().splitlines()
    text = ""
    for line in lines:
        item = __import__("json").loads(line)
        for res in item["result"]["layoutParsingResults"]:
            text += res["markdown"]["text"]

    return text