import sys
import json
import base64
from pathlib import Path
try:
    import nbformat
except:
    print(f"Warning in ipynb2md: Failed to import nbformat")
    raise f"Error in ipynb2md: Failed to import nbformat"

# -----------------------------
# JSON → Markdown（表形式）
# -----------------------------
def json_to_markdown_table(data):
    if isinstance(data, dict):
        rows = ["| key | value |", "|-----|--------|"]
        for k, v in data.items():
            rows.append(f"| {k} | {v} |")
        return "\n".join(rows)

    elif isinstance(data, list):
        rows = ["| index | value |", "|--------|--------|"]
        for i, v in enumerate(data):
            rows.append(f"| {i} | {v} |")
        return "\n".join(rows)

    else:
        return str(data)


# -----------------------------
# Notebook → Markdown
# -----------------------------
def save_image(base64_data, out_dir, prefix):
    out_dir.mkdir(exist_ok=True)
    img_bytes = base64.b64decode(base64_data)
    img_path = out_dir / f"{prefix}.png"
    img_path.write_bytes(img_bytes)
    return img_path


def convert_ipynb_to_md(ipynb_path: Path, md_path: Path):
    nb = nbformat.read(ipynb_path, as_version=4)
    out_dir = md_path.parent / (md_path.stem + "_files")

    lines = []

    for cell_index, cell in enumerate(nb.cells):
        if cell.cell_type == "markdown":
            # attachments の処理
            if "attachments" in cell:
                for name, attach in cell["attachments"].items():
                    if "image/png" in attach:
                        img_path = save_image(
                            attach["image/png"], out_dir, f"attach_{cell_index}_{name}"
                        )
                        cell.source = cell.source.replace(
                            f"attachment:{name}", str(img_path)
                        )

            lines.append(cell.source.rstrip())
            lines.append("")

        elif cell.cell_type == "code":
            lines.append("```python")
            lines.append(cell.source.rstrip())
            lines.append("```")
            lines.append("")

            # outputs の処理
            for out_index, out in enumerate(cell.get("outputs", [])):
                if out.output_type == "stream":
                    lines.append("```\n" + out.text.rstrip() + "\n```")
                    lines.append("")

                elif out.output_type in ("execute_result", "display_data"):
                    data = out.get("data", {})

                    # 画像
                    if "image/png" in data:
                        img_path = save_image(
                            data["image/png"], out_dir, f"output_{cell_index}_{out_index}"
                        )
                        lines.append(f"![output]({img_path})")
                        lines.append("")

                    # テキスト
                    if "text/plain" in data:
                        lines.append("```\n" + data["text/plain"].rstrip() + "\n```")
                        lines.append("")

                elif out.output_type == "error":
                    lines.append("```text")
                    lines.append("\n".join(out["traceback"]))
                    lines.append("```")
                    lines.append("")

        else:
            lines.append(f"<!-- Unsupported cell type: {cell.cell_type} -->")
            lines.append(cell.source.rstrip())
            lines.append("")

    md_path.write_text("\n".join(lines).rstrip() + "\n", encoding="utf-8")
    print(f"Written: {md_path}")


# -----------------------------
# JSON → Markdown
# -----------------------------
def convert_json_to_md(json_path: Path, md_path: Path):
    data = json.loads(json_path.read_text(encoding="utf-8"))
    md = json_to_markdown_table(data)
    md_path.write_text(md + "\n", encoding="utf-8")
    print(f"Written: {md_path}")


# -----------------------------
# メイン
# -----------------------------
def main():
    if len(sys.argv) < 2:
        print("Usage: python convert.py input.(ipynb|json) [output.md]")
        sys.exit(1)

    input_path = Path(sys.argv[1])
    output_path = (
        Path(sys.argv[2]) if len(sys.argv) >= 3
        else input_path.with_suffix(".md")
    )

    ext = input_path.suffix.lower()

    if ext == ".ipynb":
        convert_ipynb_to_md(input_path, output_path)
    elif ext == ".json":
        convert_json_to_md(input_path, output_path)
    else:
        print(f"Unsupported extension: {ext}")
        sys.exit(1)


if __name__ == "__main__":
    main()