import base64
import json
import os
import re
import shutil
import threading
import traceback
import urllib.error
import urllib.parse
import urllib.request
import zipfile
from datetime import datetime, timezone
from uuid import uuid4

from flask import Flask, jsonify, render_template, request, send_from_directory
from PIL import Image, ImageOps
import anthropic
try:
    from openai import OpenAI
except Exception:
    OpenAI = None

APP_ROOT = os.path.dirname(os.path.abspath(__file__))
OUTPUT_DIR = os.path.join(APP_ROOT, "outputs")
GEMINI_API_BASE = "https://generativelanguage.googleapis.com/v1beta"

app = Flask(__name__, static_folder="static", template_folder="templates")

PROMPTS = {
    "extract": """The image depicts a figure from a research paper. Please carefully analyze the figure and extract the horizontal (\"x\") and vertical (\"y\") coordinates of each individual data point from each curve. Pay extremely close attention to the axes:
- Identify the tick marks and labels on both axes, and use them to determine the numeric values for every point.
- Read the labels of both axes (the text or symbols denoting what each axis represents).

Data Reporting Requirements:
1. Separate each curve's data so that every curve has its own two columns of x and y.
2. The first row of your output must contain the axis labels (for example, \"x-axis label\" for the x-column and \"y-axis label (Curve X)\" for the y-column). Include the curve label in the y-axis header if it's available.
3. Subsequent rows must present the extracted numeric data in comma-separated values (CSV) format, with:
   - No additional words or commentary.
   - Exactly one row per data point in each curve.
   - For multiple curves, produce additional pairs of columns (two columns per curve: one for x-values, one for y-values).
4. If, for any reason, you are unable to extract the data, respond with only the word \"None\".
5. Make sure that both axis are labeled and there are individual values assigned to each tickmark, and that they increase in a logical monotonius fashion. If they are not it may be impossible to extract the data. In that case respond with only the word \"None\".

Key Instructions:
- Look at each point individually; do not assume a pattern or formula. Each data point must be read precisely from its plotted location.
- Verify the numerical axis labels carefully so the extracted data is as accurate as possible.
- Do not include any explanations or text other than the required CSV table.

Output Format Example (if you have two curves):
x-axis label,y-axis label (Curve 1),x-axis label,y-axis label (Curve 2)
0.1,0.5,0.1,0.75
0.2,0.52,0.2,0.77
...

(Or \"None\" if you cannot extract the data.)

Remember: The sole output should be either the CSV table (with all columns for the curves) or \"None\". Nothing else.
""",
    "code_fix": "The text above is an error produced by your code, please fix the code so that this error does not appear. Use only matplotlib and numpy (plus the Python standard library). Repeat the whole code and only the code so that your whole response can be directly copied and executed. Do not explain and do not say anything else, respond with just the code. Do not output the markdown code fence.",
    "compare_x": "You are provided with two images of research plots extracted from academic papers. Do these two plots have the same x-axis (horizontal)? Do they have the same ranges, labels, etc.? Answer with a single word, \"yes\" or \"no\" only.",
    "compare_y": "You are provided with two images of research plots extracted from academic papers. Do these two plots have the same y-axis (vertical)? Do they have the same ranges, labels, etc.? Answer with a single word, \"yes\" or \"no\" only.",
    "compare_number": "You are provided with two images of research plots extracted from academic papers. Do these two plots have the same number of points (for point plots)? Do the curves look like they connect the same amount of pointis (for line plots)?. Answer with a single word, \"yes\" or \"no\" only.",
    "compare_trend": "You are provided with two images of research plots extracted from academic papers. Do these sets of points or curves on these two plots represent the same trends? Do they follow the same patterns? Are points distributed in the same way? Answer with a single word, \"yes\" or \"no\" only.",
}


class GeminiClient:
    def __init__(self, api_key):
        self.api_key = api_key

    def _request(self, method, path, payload=None, params=None):
        query = dict(params or {})
        query["key"] = self.api_key
        url = f"{GEMINI_API_BASE}{path}?{urllib.parse.urlencode(query)}"
        data = None
        headers = {}
        if payload is not None:
            data = json.dumps(payload).encode("utf-8")
            headers["Content-Type"] = "application/json"
        request_obj = urllib.request.Request(url, data=data, headers=headers, method=method)

        try:
            with urllib.request.urlopen(request_obj, timeout=120) as response:
                body = response.read().decode("utf-8")
        except urllib.error.HTTPError as exc:
            body = exc.read().decode("utf-8", errors="replace")
            try:
                payload = json.loads(body)
                message = payload.get("error", {}).get("message") or body
            except Exception:
                message = body
            raise RuntimeError(message) from exc
        except urllib.error.URLError as exc:
            raise RuntimeError(str(exc.reason)) from exc

        return json.loads(body) if body else {}

    def list_models(self, page_token=""):
        params = {"pageSize": 1000}
        if page_token:
            params["pageToken"] = page_token
        return self._request("GET", "/models", params=params)

    def generate_content(self, model, contents, max_output_tokens):
        model_name = gemini_model_name(model)
        model_path = urllib.parse.quote(model_name, safe="/")
        return self._request(
            "POST",
            f"/{model_path}:generateContent",
            payload={
                "contents": contents,
                "generationConfig": {"maxOutputTokens": max_output_tokens},
            },
        )


def detect_provider(api_key):
    key = (api_key or "").strip()
    if key.startswith("sk-ant-"):
        return "anthropic"
    if key.startswith("sk-"):
        return "openai"
    return "gemini" if key else ""


def init_provider_client(api_key, provider):
    if provider == "anthropic":
        return anthropic.Anthropic(api_key=api_key)
    if provider == "openai":
        if OpenAI is None:
            raise RuntimeError("OpenAI SDK is not installed on the server.")
        return OpenAI(api_key=api_key)
    if provider == "gemini":
        return GeminiClient(api_key=api_key)
    raise ValueError("Unsupported API key format. Use an Anthropic, OpenAI, or Google Gemini API key.")


def provider_label(provider):
    if provider == "anthropic":
        return "Anthropic"
    if provider == "openai":
        return "OpenAI"
    if provider == "gemini":
        return "Google Gemini"
    return "Unknown"


def app_path(path):
    prefix = (request.script_root or "").rstrip("/")
    return f"{prefix}{path}"


@app.route("/")
def index():
    return render_template("index.html")


@app.route("/outputs/<path:filename>")
def outputs(filename):
    return send_from_directory(OUTPUT_DIR, filename)


@app.route("/api/run", methods=["POST"])
def run_extraction():
    api_key = (request.form.get("api_key") or "").strip()
    model = (request.form.get("model") or "").strip()
    image_file = request.files.get("image")
    provider = detect_provider(api_key)

    if not api_key:
        return jsonify({"error": "Missing API key."}), 400
    if not provider:
        return jsonify({"error": "Unsupported API key format. Use an Anthropic, OpenAI, or Google Gemini API key."}), 400
    if not image_file:
        return jsonify({"error": "Missing image upload."}), 400

    run_id = uuid4().hex[:10]
    run_dir = os.path.join(OUTPUT_DIR, run_id)
    os.makedirs(run_dir, exist_ok=True)

    ext = os.path.splitext(image_file.filename)[-1].lower()
    if ext not in [".png", ".jpg", ".jpeg", ".zip"]:
        return jsonify({"error": "Please upload a PNG, JPG, or ZIP file."}), 400

    try:
        client = init_provider_client(api_key, provider)
    except Exception as exc:
        return jsonify({
            "error": f"Failed to initialize {provider_label(provider)} client. Check your API key.",
            "details": str(exc),
        }), 400

    if not model:
        try:
            models = fetch_all_models(client, provider)
            model = choose_default_model(models, provider)
        except Exception as exc:
            return jsonify({
                "error": f"Failed to fetch available models from {provider_label(provider)}.",
                "details": str(exc),
            }), 500
        if not model:
            return jsonify({"error": "No available vision-capable models found for this API key."}), 400

    try:
        if ext == ".zip":
            zip_path = os.path.join(run_dir, "input.zip")
            image_file.save(zip_path)
            result = prepare_zip_run(zip_path, run_dir, run_id)
            if "error" in result:
                return jsonify({"error": result["error"]}), 400
            result["base_path"] = request.script_root or ""
            start_zip_worker(api_key, provider, model, result)
            return jsonify({
                "run_id": run_id,
                "status": "zip_pending",
                "progress_url": app_path(f"/api/progress/{run_id}"),
                "results_zip_url": app_path(f"/outputs/{run_id}/results.zip"),
            })

        input_name = f"input{ext if ext != '.jpeg' else '.jpg'}"
        input_path = os.path.join(run_dir, input_name)
        image_file.save(input_path)

        result = process_single_image(client, provider, model, input_path, run_dir)
        if result["status"] == "none":
            return jsonify({
                "run_id": run_id,
                "data_csv": result.get("data_csv", "None"),
                "status": "none",
                "input_url": app_path(f"/outputs/{run_id}/{input_name}"),
            })
        if result["status"] == "error":
            return jsonify({
                "error": "Replot code failed after retries.",
                "details": result.get("details", ""),
                "data_csv": result.get("data_csv", ""),
                "run_id": run_id,
            }), 500

        return jsonify({
            "run_id": run_id,
            "status": "ok",
            "data_csv": result["data_csv"],
            "code": result["code"],
            "validation": result["validation"],
            "validated": result["validated"],
            "input_url": app_path(f"/outputs/{run_id}/{input_name}"),
            "replot_url": app_path(f"/outputs/{run_id}/replot.png"),
            "stacked_url": app_path(f"/outputs/{run_id}/stacked.png"),
            "comparison_url": app_path(f"/outputs/{run_id}/comparison.png"),
        })
    except Exception as exc:
        return jsonify({
            "error": "Unexpected error.",
            "details": str(exc),
            "trace": traceback.format_exc(),
        }), 500


@app.route("/api/progress/<run_id>")
def progress(run_id):
    progress_path = os.path.join(OUTPUT_DIR, run_id, "progress.json")
    if not os.path.exists(progress_path):
        return jsonify({"error": "Progress not found."}), 404
    try:
        with open(progress_path, "r", encoding="utf-8") as handle:
            return jsonify(json.loads(handle.read()))
    except Exception:
        return jsonify({"error": "Failed to read progress."}), 500


@app.route("/api/models", methods=["POST"])
def list_models():
    payload = request.get_json(silent=True) or {}
    api_key = (payload.get("api_key") or request.form.get("api_key") or "").strip()
    provider = detect_provider(api_key)

    if not api_key:
        return jsonify({"error": "Missing API key."}), 400
    if not provider:
        return jsonify({"error": "Unsupported API key format. Use an Anthropic, OpenAI, or Google Gemini API key."}), 400

    try:
        client = init_provider_client(api_key, provider)
    except Exception as exc:
        return jsonify({
            "error": f"Failed to initialize {provider_label(provider)} client. Check your API key.",
            "details": str(exc),
        }), 400

    try:
        models = fetch_all_models(client, provider)
        return jsonify({
            "models": models,
            "default_model": choose_default_model(models, provider),
            "provider": provider,
        })
    except Exception as exc:
        return jsonify({
            "error": f"Failed to fetch models from {provider_label(provider)}.",
            "details": str(exc),
        }), 500


def encode_image(image_path):
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode("utf-8")


def parse_model_list_response(items):
    models = []
    seen = set()
    for item in items or []:
        if isinstance(item, dict):
            model_id = item.get("id")
            display_name = item.get("display_name") or item.get("name")
            created_at = item.get("created_at") or item.get("created")
        else:
            model_id = getattr(item, "id", None)
            display_name = getattr(item, "display_name", None) or getattr(item, "name", None)
            created_at = getattr(item, "created_at", None) or getattr(item, "created", None)

        if not model_id or model_id in seen:
            continue
        seen.add(model_id)

        created_text = ""
        if created_at is not None:
            if isinstance(created_at, (int, float)):
                try:
                    created_text = datetime.fromtimestamp(created_at, tz=timezone.utc).isoformat()
                except Exception:
                    created_text = str(created_at)
            elif hasattr(created_at, "isoformat"):
                created_text = created_at.isoformat()
            else:
                created_text = str(created_at)

        models.append({
            "id": str(model_id),
            "display_name": str(display_name or model_id),
            "created_at": created_text,
        })
    return models


def is_anthropic_vision_model(model):
    model_id = (model.get("id") or "").lower()
    if not model_id.startswith("claude"):
        return False
    return any(family in model_id for family in ["opus", "sonnet", "haiku", "fable", "mythos"])


def is_openai_vision_model(model):
    model_id = (model.get("id") or "").lower()
    if any(excluded in model_id for excluded in [
        "audio",
        "dall-e",
        "embedding",
        "gpt-image",
        "moderation",
        "realtime",
        "search",
        "transcribe",
        "tts",
        "whisper",
    ]):
        return False
    if model_id.startswith("computer-use-preview"):
        return True
    if model_id.startswith(("gpt-5", "gpt-4o", "gpt-4.1", "gpt-4.5")):
        return True
    return bool(re.match(r"^o(1|3|4)(?:-|$)", model_id))


def is_gemini_vision_model(model):
    model_id = (model.get("id") or "").lower()
    if "embedding" in model_id:
        return False
    return "gemini" in model_id


def fetch_all_anthropic_models(client):
    page = client.models.list(limit=1000)
    collected = []

    while True:
        page_data = getattr(page, "data", None) or []
        collected.extend(page_data)
        if hasattr(page, "has_next_page") and page.has_next_page():
            page = page.get_next_page()
            continue
        break

    models = parse_model_list_response(collected)
    return [model for model in models if is_anthropic_vision_model(model)]


def fetch_all_openai_models(client):
    page = client.models.list()
    collected = []

    while True:
        page_data = getattr(page, "data", None) or []
        collected.extend(page_data)
        if hasattr(page, "has_next_page") and page.has_next_page():
            page = page.get_next_page()
            continue
        break

    models = parse_model_list_response(collected)
    return [model for model in models if is_openai_vision_model(model)]


def gemini_model_name(model):
    value = (model or "").strip()
    if value.startswith("models/"):
        return value
    return f"models/{value}"


def parse_gemini_model_list_response(items):
    models = []
    seen = set()
    for item in items or []:
        model_name = item.get("name") or item.get("baseModelId")
        model_id = str(model_name or "").strip()
        if model_id.startswith("models/"):
            model_id = model_id[len("models/"):]
        methods = item.get("supportedGenerationMethods") or []
        if methods and "generateContent" not in methods:
            continue
        if not model_id or model_id in seen:
            continue
        seen.add(model_id)
        display_name = item.get("displayName") or item.get("baseModelId") or model_id
        models.append({
            "id": model_id,
            "display_name": str(display_name),
            "created_at": "",
        })
    return models


def fetch_all_gemini_models(client):
    collected = []
    page_token = ""

    while True:
        page = client.list_models(page_token=page_token)
        collected.extend(page.get("models") or [])
        page_token = page.get("nextPageToken") or ""
        if not page_token:
            break

    models = parse_gemini_model_list_response(collected)
    return [model for model in models if is_gemini_vision_model(model)]


def gemini_sort_key(model):
    # Gemini lists carry no release date: the version number tracks release
    # order, then pro > flash > flash-lite, then stable before preview.
    # "-latest" aliases float to the top.
    model_id = (model.get("id") or "").lower()
    version = re.search(r"-(\d+(?:\.\d+)?)(?:-|$)", model_id)
    if "lite" in model_id:
        tier = 1
    elif "flash" in model_id:
        tier = 2
    elif "pro" in model_id:
        tier = 3
    else:
        tier = 0
    return (
        "latest" in model_id,
        float(version.group(1)) if version else 0.0,
        tier,
        not any(tag in model_id for tag in ["preview", "exp"]),
        model_id,
    )


def sort_models_newest_first(models, provider):
    if provider == "gemini":
        return sorted(models, key=gemini_sort_key, reverse=True)
    # Anthropic and OpenAI report a creation time (UTC ISO strings sort correctly)
    return sorted(models, key=lambda model: (model.get("created_at") or "", model.get("id") or ""), reverse=True)


def fetch_all_models(client, provider):
    if provider == "anthropic":
        return sort_models_newest_first(fetch_all_anthropic_models(client), provider)
    if provider == "openai":
        return sort_models_newest_first(fetch_all_openai_models(client), provider)
    if provider == "gemini":
        return sort_models_newest_first(fetch_all_gemini_models(client), provider)
    return []


def choose_default_model(models, provider):
    # fetch_all_models returns lists sorted newest-first, so the default is
    # simply the top entry.
    return models[0]["id"] if models else ""


def media_type_for_path(image_path):
    ext = os.path.splitext(image_path)[-1].lower()
    if ext in [".jpg", ".jpeg"]:
        return "image/jpeg"
    return "image/png"


def write_progress(run_dir, data):
    progress_path = os.path.join(run_dir, "progress.json")
    try:
        with open(progress_path, "w", encoding="utf-8") as handle:
            handle.write(json.dumps(data))
    except Exception:
        pass


def process_single_image(client, provider, model, input_path, output_dir):
    os.makedirs(output_dir, exist_ok=True)
    replot_path = os.path.join(output_dir, "replot.png")
    stacked_path = os.path.join(output_dir, "stacked.png")
    comparison_path = os.path.join(output_dir, "comparison.png")
    data_path = os.path.join(output_dir, "data.csv")
    code_path = os.path.join(output_dir, "replot_code.py")
    validation_path = os.path.join(output_dir, "validation.json")

    base64_image = encode_image(input_path)
    context, data_csv = prompt_extract(client, provider, model, base64_image, media_type_for_path(input_path))

    if not data_csv or "none" in data_csv.strip().lower():
        try:
            with open(data_path, "w", encoding="utf-8") as handle:
                handle.write(data_csv.strip() if data_csv else "None")
        except Exception:
            pass
        return {"status": "none", "data_csv": data_csv.strip() if data_csv else "None"}

    code, convo = prompt_replot_code(client, provider, model, data_csv, replot_path, context)
    code, error_output = exec_replot_code(client, provider, model, code, replot_path, convo)

    if error_output:
        error_log = os.path.join(output_dir, "replot_error.txt")
        try:
            with open(error_log, "w", encoding="utf-8") as handle:
                handle.write(error_output)
        except Exception:
            pass
        return {"status": "error", "details": error_output, "data_csv": data_csv}

    try:
        with open(data_path, "w", encoding="utf-8") as handle:
            handle.write(data_csv)
        with open(code_path, "w", encoding="utf-8") as handle:
            handle.write(code)
    except Exception:
        pass

    stacked_path = stack_images_vertically(
        input_path,
        replot_path,
        stacked_path,
        border_color=None,
        border_size=0,
    )
    validation = run_comparisons(client, provider, model, stacked_path)
    validated = "yes" if not validation["wrong"] else "no"

    comparison_path = stack_images_vertically(
        input_path,
        replot_path,
        comparison_path,
        border_color=validated,
        border_size=24,
    )

    try:
        with open(validation_path, "w", encoding="utf-8") as handle:
            handle.write(json.dumps(validation, indent=2))
    except Exception:
        pass

    return {
        "status": "ok",
        "data_csv": data_csv,
        "code": code,
        "validation": validation,
        "validated": validated,
    }


def prepare_zip_run(zip_path, run_dir, run_id):
    extract_dir = os.path.join(run_dir, "zip_in")
    results_dir = os.path.join(run_dir, "results")
    os.makedirs(extract_dir, exist_ok=True)
    os.makedirs(results_dir, exist_ok=True)

    extract_zip_safe(zip_path, extract_dir)

    image_paths = []
    for root, _, files in os.walk(extract_dir):
        for name in files:
            lower = name.lower()
            if lower.endswith((".png", ".jpg", ".jpeg")):
                image_paths.append(os.path.join(root, name))

    image_paths.sort()
    if not image_paths:
        return {"error": "No PNG or JPG files found in the ZIP."}

    used_names = set()
    jobs = []
    for idx, image_path in enumerate(image_paths, 1):
        base = os.path.splitext(os.path.basename(image_path))[0]
        safe_base = "".join(ch if ch.isalnum() else "_" for ch in base).strip("_") or f"image_{idx}"
        name = safe_base
        suffix = 1
        while name in used_names:
            suffix += 1
            name = f"{safe_base}_{suffix}"
        used_names.add(name)

        item_dir = os.path.join(results_dir, name)
        os.makedirs(item_dir, exist_ok=True)
        ext = os.path.splitext(image_path)[-1].lower()
        input_name = f"input{ext if ext != '.jpeg' else '.jpg'}"
        input_path = os.path.join(item_dir, input_name)
        shutil.copy(image_path, input_path)

        jobs.append({
            "name": name,
            "input_path": input_path,
            "item_dir": item_dir,
        })

    write_progress(run_dir, {
        "status": "running",
        "total": len(jobs),
        "processed": 0,
        "current": "",
        "items": [],
    })

    return {
        "run_id": run_id,
        "run_dir": run_dir,
        "results_dir": results_dir,
        "jobs": jobs,
    }

def start_zip_worker(api_key, provider, model, payload):
    thread = threading.Thread(
        target=process_zip_async,
        args=(api_key, provider, model, payload),
        daemon=True,
    )
    thread.start()


def process_zip_async(api_key, provider, model, payload):
    run_dir = payload["run_dir"]
    results_dir = payload["results_dir"]
    jobs = payload["jobs"]

    try:
        client = init_provider_client(api_key, provider)
    except Exception as exc:
        write_progress(run_dir, {
            "status": "error",
            "error": f"Failed to initialize {provider_label(provider)} client.",
            "details": str(exc),
        })
        return

    items = []
    total = len(jobs)
    for idx, job in enumerate(jobs, 1):
        write_progress(run_dir, {
            "status": "running",
            "total": total,
            "processed": idx - 1,
            "current": job["name"],
            "items": items,
        })
        result = process_single_image(client, provider, model, job["input_path"], job["item_dir"])
        items.append({
            "name": job["name"],
            "status": result.get("status"),
            "validated": result.get("validated", ""),
            "wrong_why": result.get("validation", {}).get("wrong_why", "") if result.get("validation") else "",
        })
        write_progress(run_dir, {
            "status": "running",
            "total": total,
            "processed": idx,
            "current": job["name"],
            "items": items,
        })

    summary_path = os.path.join(results_dir, "summary.json")
    try:
        with open(summary_path, "w", encoding="utf-8") as handle:
            handle.write(json.dumps({"items": items}, indent=2))
    except Exception:
        pass

    zip_out = os.path.join(run_dir, "results.zip")
    with zipfile.ZipFile(zip_out, "w", compression=zipfile.ZIP_DEFLATED) as zf:
        for root, _, files in os.walk(results_dir):
            for file_name in files:
                full_path = os.path.join(root, file_name)
                rel_path = os.path.relpath(full_path, results_dir)
                zf.write(full_path, arcname=rel_path)

    write_progress(run_dir, {
        "status": "done",
        "total": total,
        "processed": total,
        "current": "",
        "items": items,
        "results_zip_url": f"{(payload.get('base_path') or '').rstrip('/')}/outputs/{payload['run_id']}/results.zip",
    })


def extract_zip_safe(zip_path, dest_dir):
    with zipfile.ZipFile(zip_path, "r") as zf:
        for member in zf.infolist():
            if member.filename.endswith("/"):
                continue
            extracted_path = os.path.normpath(os.path.join(dest_dir, member.filename))
            if not extracted_path.startswith(os.path.abspath(dest_dir) + os.sep):
                continue
            os.makedirs(os.path.dirname(extracted_path), exist_ok=True)
            with zf.open(member) as source, open(extracted_path, "wb") as target:
                shutil.copyfileobj(source, target)


def extract_openai_text(response):
    if isinstance(response, dict):
        output_text = response.get("output_text")
        if output_text:
            return str(output_text)
        outputs = response.get("output") or []
    else:
        output_text = getattr(response, "output_text", None)
        outputs = getattr(response, "output", None) or []
    if output_text:
        return output_text

    parts = []
    for item in outputs:
        content = item.get("content") if isinstance(item, dict) else getattr(item, "content", None)
        for block in content or []:
            block_type = block.get("type") if isinstance(block, dict) else getattr(block, "type", None)
            if block_type not in ["output_text", "text"]:
                continue
            text = block.get("text") if isinstance(block, dict) else getattr(block, "text", None)
            if not text and not isinstance(block, dict):
                text = getattr(block, "value", None)
            if text:
                parts.append(str(text))
    return "\n".join(parts).strip()


def extract_gemini_text(response):
    parts = []
    for candidate in response.get("candidates") or []:
        content = candidate.get("content") or {}
        for part in content.get("parts") or []:
            text = part.get("text")
            if text:
                parts.append(str(text))
    return "\n".join(parts).strip()


def prompt_extract(client, provider, model, base64_image, media_type):
    if provider == "anthropic":
        messages = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "image",
                        "source": {
                            "type": "base64",
                            "media_type": media_type,
                            "data": base64_image,
                        },
                    },
                    {"type": "text", "text": PROMPTS["extract"]},
                ],
            }
        ]
        response = client.messages.create(
            model=model,
            max_tokens=2048,
            #temperature=0,
            messages=messages,
        )
        return messages, response.content[0].text

    if provider == "gemini":
        messages = [
            {
                "role": "user",
                "parts": [
                    {
                        "inline_data": {
                            "mime_type": media_type,
                            "data": base64_image,
                        },
                    },
                    {"text": PROMPTS["extract"]},
                ],
            }
        ]
        response = client.generate_content(model, messages, max_output_tokens=2048)
        return messages, extract_gemini_text(response)

    image_data_url = f"data:{media_type};base64,{base64_image}"
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "input_image", "image_url": image_data_url},
                {"type": "input_text", "text": PROMPTS["extract"]},
            ],
        }
    ]
    response = client.responses.create(
        model=model,
        input=messages,
        max_output_tokens=2048,
    )
    return messages, extract_openai_text(response)


def prompt_replot_code(client, provider, model, data_csv, replot_path, messages):
    prompt = (
        "Please analyze the figure and create a python code that will reproduce the plot exactly, "
        "including colors, line types, point shapes, axis labels, axis ranges, etc. "
        f"Save the plot as a file \"{replot_path}\" only and do not show it. "
        "Use only matplotlib and numpy (plus the Python standard library). "
        "Respond with the code only so that it can be directly copied and executed, do not output the markdown code fence.\n\n"
        f"Use the following data on the plot: {data_csv}"
    )
    if provider == "anthropic":
        convo = list(messages)
        convo.append({"role": "assistant", "content": data_csv})
        convo.append({"role": "user", "content": prompt})
        response = client.messages.create(
            model=model,
            max_tokens=2048,
            #temperature=0,
            messages=convo,
        )
        return response.content[0].text, convo

    if provider == "gemini":
        convo = list(messages)
        convo.append({"role": "model", "parts": [{"text": data_csv}]})
        convo.append({"role": "user", "parts": [{"text": prompt}]})
        response = client.generate_content(model, convo, max_output_tokens=2048)
        return extract_gemini_text(response), convo

    convo = list(messages)
    convo.append({
        "role": "user",
        "content": [
            {
                "type": "input_text",
                "text": f"Use the following extracted CSV data:\n{data_csv}\n\n{prompt}",
            }
        ],
    })
    response = client.responses.create(
        model=model,
        input=convo,
        max_output_tokens=2048,
    )
    return extract_openai_text(response), convo


def exec_replot_code(client, provider, model, code, replot_path, convo=None):
    error_output = None

    def strip_code_fences(text):
        if not text:
            return text
        cleaned = text.strip()
        if cleaned.startswith("```"):
            lines = cleaned.splitlines()
            lines = lines[1:]
            if lines and lines[-1].strip().startswith("```"):
                lines = lines[:-1]
            cleaned = "\n".join(lines).strip()
        return cleaned

    def run_code(script):
        exec_globals = {"__name__": "__main__"}
        try:
            import matplotlib
            matplotlib.use("Agg")
            import matplotlib.pyplot as plt
            plt.show = lambda *args, **kwargs: None
            exec_globals["plt"] = plt
        except Exception:
            pass
        exec(script, exec_globals)

    try:
        code = strip_code_fences(code)
        run_code(code)
    except Exception:
        error_output = traceback.format_exc()

    attempts = 0
    while error_output and attempts < 2:
        attempts += 1
        fix_prompt = error_output + "\n\n" + PROMPTS["code_fix"]
        if provider == "anthropic":
            fix_messages = [
                {"role": "assistant", "content": code},
                {"role": "user", "content": fix_prompt},
            ]
            if convo:
                fix_messages = list(convo) + fix_messages
            response = client.messages.create(
                model=model,
                max_tokens=2048,
                #temperature=0,
                messages=fix_messages,
            )
            code = strip_code_fences(response.content[0].text)
        elif provider == "gemini":
            fix_messages = []
            if convo:
                fix_messages = list(convo)
            fix_messages.extend([
                {"role": "model", "parts": [{"text": code}]},
                {"role": "user", "parts": [{"text": fix_prompt}]},
            ])
            response = client.generate_content(model, fix_messages, max_output_tokens=2048)
            code = strip_code_fences(extract_gemini_text(response))
        else:
            fix_messages = []
            if convo:
                fix_messages = list(convo) + fix_messages
            fix_messages.append(
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "input_text",
                            "text": f"Current code:\n{code}\n\nExecution error:\n{fix_prompt}",
                        }
                    ],
                }
            )
            response = client.responses.create(
                model=model,
                input=fix_messages,
                max_output_tokens=2048,
            )
            code = strip_code_fences(extract_openai_text(response))
        try:
            run_code(code)
            error_output = None
        except Exception:
            error_output = traceback.format_exc()

    return code, error_output


def run_comparisons(client, provider, model, stacked_path):
    base64_image = encode_image(stacked_path)
    media_type = media_type_for_path(stacked_path)
    results = {}
    wrong = False
    wrong_why = []

    for key, label in [
        ("compare_x", "X"),
        ("compare_y", "Y"),
        ("compare_number", "N"),
        ("compare_trend", "T"),
    ]:
        if provider == "anthropic":
            response = client.messages.create(
                model=model,
                max_tokens=32,
                #temperature=0,
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "image",
                                "source": {
                                    "type": "base64",
                                    "media_type": media_type,
                                    "data": base64_image,
                                },
                            },
                            {"type": "text", "text": PROMPTS[key]},
                        ],
                    }
                ],
            )
            answer = response.content[0].text.strip().lower()
        elif provider == "gemini":
            response = client.generate_content(
                model,
                [
                    {
                        "role": "user",
                        "parts": [
                            {
                                "inline_data": {
                                    "mime_type": media_type,
                                    "data": base64_image,
                                },
                            },
                            {"text": PROMPTS[key]},
                        ],
                    }
                ],
                max_output_tokens=32,
            )
            answer = extract_gemini_text(response).strip().lower()
        else:
            image_data_url = f"data:{media_type};base64,{base64_image}"
            response = client.responses.create(
                model=model,
                input=[
                    {
                        "role": "user",
                        "content": [
                            {"type": "input_image", "image_url": image_data_url},
                            {"type": "input_text", "text": PROMPTS[key]},
                        ],
                    }
                ],
                max_output_tokens=32,
            )
            answer = extract_openai_text(response).strip().lower()
        results[key] = answer
        if answer.startswith("no"):
            wrong = True
            wrong_why.append(label)

    return {
        "answers": results,
        "wrong": wrong,
        "wrong_why": "; ".join(wrong_why) if wrong_why else "",
    }


def stack_images_vertically(image1_path, image2_path, output_path, border_color=None, border_size=30):
    img1 = Image.open(image1_path).convert("RGB")
    img2 = Image.open(image2_path).convert("RGB")

    width = max(img1.width, img2.width)
    img1_resized = img1.resize((width, int(img1.height * width / img1.width)), Image.LANCZOS)
    img2_resized = img2.resize((width, int(img2.height * width / img2.width)), Image.LANCZOS)

    combined_height = img1_resized.height + img2_resized.height
    combined = Image.new("RGB", (width, combined_height), (0, 0, 0))
    combined.paste(img1_resized, (0, 0))
    combined.paste(img2_resized, (0, img1_resized.height))

    if border_color:
        if "yes" in border_color.lower():
            color = (16, 185, 129)
        elif "no" in border_color.lower():
            color = (239, 68, 68)
        else:
            color = (99, 102, 241)
        combined = ImageOps.expand(combined, border=border_size, fill=color)

    combined.save(output_path)
    return output_path


if __name__ == "__main__":
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    app.run(debug=True, port=5173)