#!/usr/bin/env python3
"""
itab.py - Convert any text table from stdin → beautiful PNG (or clipboard)
"""

import sys
import argparse
import io
import pandas as pd
from PIL import Image, ImageDraw, ImageFont
import pyperclip

# -------------------------- Themes --------------------------
THEMES = {
    "default":   {"bg": (255,255,255), "header_bg": (40,110,180),   "text": (0,0,0),     "header_text": (255,255,255), "grid": (200,200,200), "alternating": (240,248,255)},
    "dark":      {"bg": (30,30,30),     "header_bg": (0,120,215),    "text": (220,220,220), "header_text": (255,255,255), "grid": (70,70,70),   "alternating": (45,45,45)},
    "excel":     {"bg": (255,255,255), "header_bg": (0,120,70),     "text": (0,0,0),     "header_text": (255,255,255), "grid": (180,180,180), "alternating": (226,239,218)},
    "minimal":   {"bg": (255,255,255), "header_bg": (255,255,255),  "text": (50,50,50),  "header_text": (30,30,30),    "grid": (220,220,220), "alternating": (250,250,250), "header_underline": (100,100,100)},
    "ocean":     {"bg": (240,248,255), "header_bg": (0,102,204),    "text": (0,0,0),     "header_text": (255,255,255), "grid": (180,200,230), "alternating": (220,235,255)},
}

def get_font(size=20):
    candidates = [
        "arial.ttf",
        "DejaVuSans.ttf",
        "/System/Library/Fonts/Arial.ttf",      # macOS
        "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",  # Linux
    ]
    for path in candidates:
        try:
            return ImageFont.truetype(path, size)
        except OSError:
            continue
    return ImageFont.load_default()

def text_size(draw, text, font):
    bbox = draw.textbbox((0, 0), text, font=font)
    return bbox[2] - bbox[0], bbox[3] - bbox[1]

def create_table_image(df, theme_name="default", cell_padding=15, row_height=45):
    theme = THEMES.get(theme_name.lower(), THEMES["default"])

    header_font = get_font(22)
    cell_font    = get_font(20)

    # Measure column widths using textbbox (Pillow 8+ compatible)
    draw_temp = ImageDraw.Draw(Image.new("RGB", (1,1)))
    col_widths = []
    for col in df.columns:
        header_w, _ = text_size(draw_temp, str(col), header_font)
        col_data_w = max([text_size(draw_temp, str(row[col]), cell_font)[0] 
                          for _, row in df.iterrows()] + [0])
        col_widths.append(max(header_w, col_data_w) + cell_padding * 2)

    img_w = sum(col_widths) + cell_padding * 2
    img_h = row_height * (len(df) + 1) + cell_padding * 2

    img = Image.new("RGB", (img_w, img_h), theme["bg"])
    draw = ImageDraw.Draw(img)

    y = cell_padding

    # Header row
    x = cell_padding
    for i, col in enumerate(df.columns):
        draw.rectangle([x, y, x + col_widths[i], y + row_height], fill=theme["header_bg"])
        tw, th = text_size(draw, str(col), header_font)
        draw.text((x + (col_widths[i] - tw) // 2, y + (row_height - th) // 2),
                  str(col), fill=theme["header_text"], font=header_font)
        x += col_widths[i]

    if theme.get("header_underline"):
        draw.line([cell_padding, y + row_height, img_w - cell_padding, y + row_height],
                  fill=theme["header_underline"], width=3)

    # Data rows
    for row_idx, (_, row) in enumerate(df.iterrows()):
        y += row_height
        bg = theme["alternating"] if row_idx % 2 == 0 else theme["bg"]
        draw.rectangle([cell_padding, y, img_w - cell_padding, y + row_height], fill=bg)

        x = cell_padding
        for i, value in enumerate(row):
            text = str(value) if pd.notna(value) else ""
            tw, th = text_size(draw, text, cell_font)
            draw.text((x + (col_widths[i] - tw) // 2, y + (row_height - th) // 2),
                      text, fill=theme["text"], font=cell_font)
            x += col_widths[i]

    # Grid lines
    for i in range(len(df) + 2):
        yy = cell_padding + i * row_height
        draw.line([cell_padding, yy, img_w - cell_padding, yy], fill=theme["grid"])
    for w in [cell_padding] + [cell_padding + sum(col_widths[:i+1]) for i in range(len(col_widths))]:
        draw.line([w, cell_padding, w, img_h - cell_padding], fill=theme["grid"])

    return img

# -------------------------- Main --------------------------
def main():
    parser = argparse.ArgumentParser(description="Text table from stdin → PNG image (or clipboard)")
    parser.add_argument("-o", "--output", help="Output PNG file. If omitted → clipboard")
    parser.add_argument("-t", "--theme", default="default",
                        help=f"Themes: {', '.join(THEMES)} (default: default)")
    parser.add_argument("--sep", help="Force delimiter (comma, tab, etc.)")
    args = parser.parse_args()

    raw = sys.stdin.read().strip()
    if not raw:
        print("No input received", file=sys.stderr)
        sys.exit(1)

    # Try to parse
    df = None
    if args.sep:
        sep = args.sep
    else:
        snippet = raw[:1000]
        sep = '\t' if '\t' in snippet else ',' if ',' in snippet else None

    # 1. CSV / TSV
    if sep:
        try:
            df = pd.read_csv(io.StringIO(raw), sep=sep, engine="python")
        except:
            pass

    # 2. Space-separated (like ps, ls, etc.)
    if df is None:
        try:
            df = pd.read_table(io.StringIO(raw), sep=r'\s+', engine="python")
        except:
            pass

    # 3. Markdown table fallback
    if df is None:
        try:
            lines = raw.splitlines()
            content_lines = [l for l in lines if l.strip() and not l.strip().startswith('|---')]
            cleaned = '\n'.join(content_lines)
            df = pd.read_csv(io.StringIO(cleaned), sep='|', engine='python')
            df = df.dropna(axis=1, how='all')  # remove empty columns from leading/trailing |
            df.columns = [c.strip() for c in df.columns]
            for col in df.columns:
                df[col] = df[col].astype(str).str.strip()
        except:
            pass

    if df is None or df.empty:
        print("Could not parse input as a table", file=sys.stderr)
        sys.exit(1)

    img = create_table_image(df, theme_name=args.theme)

    if args.output:
        img.save(args.output, "PNG")
        print(f"Saved → {args.output}")
    else:
        # Copy to clipboard (cross-platform)
        output = io.BytesIO()
        img.save(output, format="PNG")
        img_bytes = output.getvalue()

        if hasattr(pyperclip, "setcb"):  # newer pyperclip supports images directly
            try:
                pyperclip.setcb(img_bytes, type="image/png")
                print("Image copied to clipboard!")
                return
            except:
                pass

        # Fallback methods
        import subprocess, platform
        system = platform.system()
        if system == "Darwin":  # macOS
            subprocess.run("pbcopy", input=img_bytes)
        elif system == "Linux":
            subprocess.run(["xclip", "-selection", "clipboard", "-t", "image/png"], input=img_bytes)
        elif system == "Windows":
            from PIL import ImageGrab
            img_pil = Image.open(io.BytesIO(img_bytes))
            ImageGrab.putclipboard(img_pil)

        print("Image copied to clipboard!")

if __name__ == "__main__":
    main()
