import json
import mimetypes
import os
import time
import urllib.request
import uuid


API_BASE = "https://retoucher.hz.labs.retouch4.me/api/v1/retoucher"
SOURCE_FILE = os.path.join(os.path.dirname(__file__), "..", "DSC_6229_Sample.jpg")
OUTPUT_FILE = os.path.join(os.path.dirname(__file__), "..", "retouch_result_python.jpg")
TOKEN = os.environ.get("RETOUCH_TOKEN", "<your_token_here>")

PAYLOAD = {
    "mode": "professional",
    "tasks": [
        {"Plugin": "Heal", "Scale": 0, "Alpha1": 1.0},
        {"Plugin": "Fabric", "Scale": 0, "Alpha1": 0.39},
        {"Plugin": "Eye Vessels", "Scale": 0, "Alpha1": 1.0},
        {"Plugin": "Eye Brilliance", "Scale": 0, "Alpha1": 0.5},
        {"Plugin": "White Teeth", "Scale": 0, "Alpha1": 0.25, "Alpha2": 0.25},
        {"Plugin": "Dodge Burn", "Scale": 2, "Alpha1": 1.0, "Alpha2": 0.2},
        {"Plugin": "Skin Tone", "Scale": 0, "Alpha1": 1.0, "Alpha2": 1.0},
        {"Plugin": "Portrait Volumes", "Scale": 0, "Alpha1": 0.5},
    ],
}


def main():
    if TOKEN == "<your_token_here>":
        raise SystemExit("Set RETOUCH_TOKEN before running this script.")

    job_id = start_job()
    print(f"Created job: {job_id}")

    wait_for_completion(job_id)
    download_result(job_id)
    print(f"Saved output to {OUTPUT_FILE}")


def start_job():
    boundary = f"----RetouchBoundary{uuid.uuid4().hex}"
    body = build_multipart_body(
        boundary,
        [
            ("token", TOKEN),
            ("payload", json.dumps(PAYLOAD, separators=(",", ":"))),
        ],
        [
            ("file", SOURCE_FILE),
        ],
    )

    request = urllib.request.Request(
        f"{API_BASE}/start",
        data=body,
        method="POST",
        headers={"Content-Type": f"multipart/form-data; boundary={boundary}"},
    )

    with urllib.request.urlopen(request) as response:
        data = json.loads(response.read().decode("utf-8"))

    return data["id"]


def wait_for_completion(job_id):
    while True:
        with urllib.request.urlopen(f"{API_BASE}/status/{job_id}") as response:
            data = json.loads(response.read().decode("utf-8"))

        state = data.get("state", "")
        print(f"State: {state}, progress: {data.get('progress', 'n/a')}")

        if state == "completed":
            return

        if state == "failed":
            raise RuntimeError(f"Job failed: {data.get('reason', 'unknown')}")

        time.sleep(5)


def download_result(job_id):
    with urllib.request.urlopen(f"{API_BASE}/getFile/{job_id}") as response:
        content = response.read()

    with open(OUTPUT_FILE, "wb") as f:
        f.write(content)


def build_multipart_body(boundary, fields, files):
    chunks = []

    for name, value in fields:
        chunks.extend(
            [
                f"--{boundary}\r\n".encode(),
                f'Content-Disposition: form-data; name="{name}"\r\n\r\n'.encode(),
                str(value).encode(),
                b"\r\n",
            ]
        )

    for field_name, file_path in files:
        filename = os.path.basename(file_path)
        content_type = mimetypes.guess_type(filename)[0] or "application/octet-stream"
        with open(file_path, "rb") as f:
            file_content = f.read()

        chunks.extend(
            [
                f"--{boundary}\r\n".encode(),
                (
                    f'Content-Disposition: form-data; name="{field_name}"; '
                    f'filename="{filename}"\r\n'
                ).encode(),
                f"Content-Type: {content_type}\r\n\r\n".encode(),
                file_content,
                b"\r\n",
            ]
        )

    chunks.append(f"--{boundary}--\r\n".encode())
    return b"".join(chunks)


if __name__ == "__main__":
    main()