#!/usr/bin/env python3
"""
pdf2md - Convert PDF files to Markdown using Google Gemini API

A command-line tool that uses the Gemini API to perform OCR and convert
PDF documents into clean, semantic Markdown format.

Usage:
    pdf2md <input.pdf> [output.md]
    python3 pdf2md.py <input.pdf> [output.md]

Environment Variables:
    MY_GEMINI_API_KEY: Your Google Gemini API key (required)

Example:
    export MY_GEMINI_API_KEY="your_api_key_here"
    pdf2md document.pdf output.md
"""

import sys
import os
import base64
import json
import urllib.request
import urllib.error

# ==========================================
# configuration
# ==========================================
# set the api key before running requests.

# you can set the environment variable in your terminal like this:
# export MY_GEMINI_API_KEY="your_api_key_here"

API_KEY = os.getenv("MY_GEMINI_API_KEY")


# use an available gemini model variant.
MODEL = "gemini-2.5-flash"
# ==========================================


def main():
    # validate cli arguments early to avoid unnecessary work.
    if len(sys.argv) < 2:
        print("Error: You must provide a PDF file as an argument.")
        print(f"Usage: {sys.argv[0]} <input.pdf> [output_name.md]")
        sys.exit(1)

    input_path = sys.argv[1]

    # fail fast when the input path is invalid.
    if not os.path.exists(input_path):
        print(f"Error: File '{input_path}' does not exist.")
        sys.exit(1)

    # keep explicit output names while normalizing extension.
    if len(sys.argv) > 2:
        output_path = sys.argv[2]
        if not output_path.endswith(".md"):
            output_path += ".md"
    else:
        filename_without_ext = os.path.splitext(os.path.basename(input_path))[0]
        output_path = f"{filename_without_ext}.md"

    print(f"Processing: {input_path}")
    print(f"Encoding file and contacting Google ({MODEL})...")

    try:
        # encode pdf bytes for the inline payload required by the api.
        with open(input_path, "rb") as pdf_file:
            pdf_data = base64.b64encode(pdf_file.read()).decode("utf-8")

        # build request payload and endpoint for content generation.
        url = f"https://generativelanguage.googleapis.com/v1beta/models/{MODEL}:generateContent?key={API_KEY}"

        headers = {"Content-Type": "application/json"}

        prompt_text = (
            "You are a highly advanced OCR and document conversion engine. "
            "Your task is to transcribe the attached PDF into clean, semantic Markdown.\n\n"
            "CRITICAL RULES:\n"
            "1. **Language:** Maintain the original language of the document exactly. Do not translate.\n"
            "2. **Structure:** Use correct Markdown headers (#, ##, ###) based on font size and hierarchy. "
            "Detect and format tables using Markdown table syntax. Use standard lists (- or 1.) for bullet points.\n"
            "3. **Math & Code:** If the document contains mathematical formulas, render them using LaTeX format (e.g., $E=mc^2$). "
            "If it contains programming code, use fenced code blocks (```python, etc).\n"
            "4. **Images:** \n"
            "   - If an image contains text/charts/graphs: Transcribe the data or text within it fully.\n"
            "   - If it is a photo/illustration: Insert a placeholder > **[Image: detailed description of the visual content]**.\n"
            "5. **Cleanliness:** Output ONLY the raw Markdown. Do not include 'Here is the file', 'I have converted...', or markdown code fences (```markdown) around the whole output."
        )

        data = {
            "contents": [
                {
                    "parts": [
                        {"text": prompt_text},
                        {
                            "inline_data": {
                                "mime_type": "application/pdf",
                                "data": pdf_data,
                            }
                        },
                    ]
                }
            ]
        }

        json_data = json.dumps(data).encode("utf-8")

        # execute request and parse the json response.
        req = urllib.request.Request(
            url, data=json_data, headers=headers, method="POST"
        )

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

        # extract the first text candidate expected in successful responses.
        try:
            markdown_content = result["candidates"][0]["content"]["parts"][0]["text"]
        except (KeyError, IndexError):
            print("Error: The model did not return text content.")
            if "candidates" in result and len(result["candidates"]) > 0:
                print(
                    f"Finish reason: {result['candidates'][0].get('finishReason', 'Unknown')}"
                )
            sys.exit(1)

        # write markdown only after extraction succeeds.
        with open(output_path, "w", encoding="utf-8") as f:
            f.write(markdown_content)

        print(f"Success! File saved to: {output_path}")

    except urllib.error.HTTPError as e:
        error_message = e.read().decode("utf-8")
        print(f"API error (code {e.code}):")
        try:
            err_json = json.loads(error_message)
            print(f"Message: {err_json['error']['message']}")
        except:
            print(error_message)
        sys.exit(1)

    except Exception as e:
        print(f"Unexpected error: {str(e)}")
        sys.exit(1)


if __name__ == "__main__":
    main()
