import io
import tempfile
import requests
import base64
import socket
from flask import Flask, render_template, request, send_from_directory, session
from PIL import Image, ImageDraw, ImageFont
import os
import re
import argparse
import json

app = Flask(__name__)
app.secret_key = 'TpfxYMxUJxJMtCCTraBSg1rbd6NLz38JTmIfpBLsotcI47EqXU'

FONT_PATH = "./fonts/ubuntu/Ubuntu-Regular.ttf"
FONT_PATH_BOLD = "./fonts/ubuntu/Ubuntu-Bold.ttf"
FONT_PATH_OBLIQUE = "./fonts/ubuntu/Ubuntu-Italic.ttf"
FONT_PATH_BOLDITALIC = "./fonts/ubuntu/Ubuntu-BoldItalic.ttf"
FONT_SIZE = 24
HEADER_SIZE_1 = 56
HEADER_SIZE_2 = 34
HEADER_SIZE_3 = 28
BANNER_FONT_SIZE = 300
IMAGE_WIDTH = 384
BULLET_CHAR = "• "
MIN_LINES = 86
BLEHD_SOCKET = os.environ.get("BLEHD_SOCKET", "/run/bleh/blehd.sock")

with open('static/translations.json', 'r', encoding='utf-8') as f:
    TRANSLATIONS = json.load(f)


def _dither_none(gray_img, levels):
    if levels == 2:
        return gray_img.point(lambda p: 0 if p < 128 else 255, mode="L")
    step = 255 / (levels - 1)
    return gray_img.point(lambda p: int(round(p / step) * step), mode="L")


def _dither_ordered(gray_img, matrix, levels, strength=1.0):
    src = gray_img.load()
    w, h = gray_img.size
    out = Image.new("L", (w, h), 255)
    dst = out.load()
    n = len(matrix)
    step = 255 / (levels - 1)
    for y in range(h):
        for x in range(w):
            t = (matrix[y % n][x % n] / (n * n) - 0.5) * step * strength
            v = max(0, min(255, int(src[x, y] + t)))
            q = int(round(v / step) * step)
            dst[x, y] = q
    return out


def _dither_error_diffusion(gray_img, levels, matrix, divisor):
    w, h = gray_img.size
    buf = [[float(gray_img.getpixel((x, y))) for x in range(w)] for y in range(h)]
    step = 255 / (levels - 1)
    for y in range(h):
        for x in range(w):
            old = buf[y][x]
            new = round(old / step) * step
            new = max(0, min(255, new))
            err = old - new
            buf[y][x] = new
            for dx, dy, weight in matrix:
                nx = x + dx
                ny = y + dy
                if 0 <= nx < w and 0 <= ny < h:
                    buf[ny][nx] += err * (weight / divisor)
    out = Image.new("L", (w, h), 255)
    for y in range(h):
        for x in range(w):
            out.putpixel((x, y), int(max(0, min(255, round(buf[y][x])))))
    return out


def _build_bayer(n):
    m = [[0, 2], [3, 1]]
    size = 2
    while size < n:
        prev = m
        size *= 2
        m = [[0] * size for _ in range(size)]
        for y in range(size // 2):
            for x in range(size // 2):
                v = prev[y][x] * 4
                m[y][x] = v + 0
                m[y][x + size // 2] = v + 2
                m[y + size // 2][x] = v + 3
                m[y + size // 2][x + size // 2] = v + 1
    return m


def apply_dither(gray_img, dither_type, levels):
    if dither_type == "none":
        return _dither_none(gray_img, levels)
    if dither_type == "floyd":
        # x+1,y ; x-1,y+1 ; x,y+1 ; x+1,y+1
        fs = [(1, 0, 7), (-1, 1, 3), (0, 1, 5), (1, 1, 1)]
        return _dither_error_diffusion(gray_img, levels, fs, 16)
    if dither_type == "atkinson":
        at = [(1, 0, 1), (2, 0, 1), (-1, 1, 1), (0, 1, 1), (1, 1, 1), (0, 2, 1)]
        return _dither_error_diffusion(gray_img, levels, at, 8)
    if dither_type == "jjn":
        # Jarvis-Judice-Ninke kernel centered at current pixel
        jjn = [
            (1, 0, 7), (2, 0, 5),
            (-2, 1, 3), (-1, 1, 5), (0, 1, 7), (1, 1, 5), (2, 1, 3),
            (-2, 2, 1), (-1, 2, 3), (0, 2, 5), (1, 2, 3), (2, 2, 1),
        ]
        return _dither_error_diffusion(gray_img, levels, jjn, 48)
    if dither_type == "bayer2x2":
        return _dither_ordered(gray_img, _build_bayer(2), levels, strength=1.0 if levels == 2 else 0.2)
    if dither_type == "bayer4x4":
        return _dither_ordered(gray_img, _build_bayer(4), levels, strength=1.0 if levels == 2 else 0.2)
    if dither_type == "bayer8x8":
        return _dither_ordered(gray_img, _build_bayer(8), levels, strength=1.0 if levels == 2 else 0.2)
    if dither_type == "bayer16x16":
        return _dither_ordered(gray_img, _build_bayer(16), levels, strength=1.0 if levels == 2 else 0.2)
    return _dither_none(gray_img, levels)


def preprocess_for_mode(img, dithering, mode):
    img = remove_alpha(img)
    if img.width <= 0 or img.height <= 0:
        raise ValueError("Invalid image dimensions")

    ratio = img.width / img.height
    height = int(IMAGE_WIDTH / ratio)
    if height <= 0:
        raise ValueError("Computed invalid image height")

    img = img.resize((IMAGE_WIDTH, height), Image.Resampling.LANCZOS).convert("L")
    levels = 2 if mode == "1bpp" else 16
    out = apply_dither(img, dithering, levels)
    return out


def pack_pixels(img, mode):
    width, height = img.size
    if width != IMAGE_WIDTH:
        img = img.resize((IMAGE_WIDTH, height), Image.Resampling.LANCZOS)
        width = IMAGE_WIDTH

    if mode == "1bpp":
        packed = bytearray((width * height) // 8)
        px = img.load()
        for y in range(height):
            for x in range(width):
                if px[x, y] < 128:
                    idx = (y * width + x) // 8
                    packed[idx] |= 1 << (x % 8)
        return bytes(packed), height

    # 4bpp
    packed = bytearray((width * height) // 2)
    px = img.load()
    for y in range(height):
        for x in range(width):
            gray = px[x, y]
            level = (255 - gray) >> 4
            idx = (y * width + x) >> 1
            shift = ((x & 1) ^ 1) << 2
            packed[idx] |= (level & 0x0F) << shift
    return bytes(packed), height


def image_from_url(url, dithering, mode):
    resp = requests.get(url, stream=True, timeout=20)
    resp.raise_for_status()
    img = Image.open(io.BytesIO(resp.content))
    return preprocess_for_mode(img, dithering, mode)


def image_from_bytes(image_bytes, dithering, mode):
    img = Image.open(io.BytesIO(image_bytes))
    return preprocess_for_mode(img, dithering, mode)


def pad_image_to_min_lines(img):
    if img.height >= MIN_LINES:
        return img
    out = Image.new("L", (img.width, MIN_LINES), 255)
    out.paste(img, (0, 0))
    return out


def blehd_call(method, params=None, socket_path=BLEHD_SOCKET, timeout=90):
    req = {"id": "1", "method": method, "params": params or {}}
    data = (json.dumps(req, ensure_ascii=False) + "\n").encode("utf-8")

    with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as s:
        s.settimeout(timeout)
        s.connect(socket_path)
        s.sendall(data)
        f = s.makefile("rb")
        line = f.readline()

    if not line:
        raise RuntimeError("No response from blehd")
    resp = json.loads(line.decode("utf-8"))
    if not resp.get("ok"):
        err = resp.get("error") or {}
        raise RuntimeError(err.get("message", "Unknown blehd error"))
    return resp.get("result")

def remove_uploaded_img():
    path = session.pop('uploaded_img_path', None)
    if path and os.path.exists(path):
        try:
            os.remove(path)
        except Exception:
            pass

def remove_alpha(img):
    if img.mode in ("RGBA", "LA"):
        background = Image.new("RGBA", img.size, (255, 255, 255, 255))  # White background
        img = Image.alpha_composite(background, img.convert("RGBA"))
        return img.convert("RGB")
    elif img.mode == "P" and 'transparency' in img.info:
        img = img.convert("RGBA")
        background = Image.new("RGBA", img.size, (255, 255, 255, 255))
        img = Image.alpha_composite(background, img)
        return img.convert("RGB")
    else:
        return img

def bleh_image_from_url(url, dithering, mode):
    return image_from_url(url, dithering, mode)


def bleh_image_from_bytes(image_bytes, dithering, mode):
    return image_from_bytes(image_bytes, dithering, mode)

def rotate_image_bytes(image_bytes, rotation):
    if rotation == 0:
        return image_bytes
    img = Image.open(io.BytesIO(image_bytes))
    img = img.rotate(-rotation, expand=True)  # PIL rotates counterclockwise, negative for clockwise
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return buf.getvalue()

def parse_line(line):
  if re.match(r"!\(img\)", line) or re.match(r"!\[.*\]\(img\)", line):
      return ('userimage', None)
  image_match = re.match(r"^!\[(.*?)\]\((.+?)\)", line)
  if image_match:
      alt = image_match.group(1)
      url = image_match.group(2)
      return ('image', url, alt)
  if re.match(r'^\s*---\s*$', line):
      return ('hr', [])
  header_match = re.match(r"^(#{1,3}) +(.*)", line)
  if header_match:
      header_level = len(header_match.group(1))
      line = header_match.group(2)
      return ('header', header_level, parse_segments(line))
  bullet_match = re.match(r"^\s*([-*\u2022]) +(.*)", line)
  if bullet_match:
      return ('bullet', parse_segments(bullet_match.group(2)))
  ordered_match = re.match(r"^\s*(\d+)\. +(.*)", line)
  if ordered_match:
      return ('ordered', int(ordered_match.group(1)), parse_segments(ordered_match.group(2)))
  return ('text', parse_segments(line))

def parse_segments(line):
    # Handle escaped asterisks: replace them with a placeholder
    line = line.replace(r'\*', '\x07')

    # Apply formatting
    bi = re.findall(r"\*\*\*(.+?)\*\*\*", line)
    for x in bi:
        line = line.replace(f"***{x}***", f"\x01{x}\x02")
    b = re.findall(r"\*\*(.+?)\*\*", line)
    for x in b:
        line = line.replace(f"**{x}**", f"\x03{x}\x04")
    i = re.findall(r"\*(.+?)\*", line)
    for x in i:
        line = line.replace(f"*{x}*", f"\x05{x}\x06")

    # Split into styled segments
    segments = []
    i = 0
    while i < len(line):
        if line[i] == '\x01':
            i += 1
            start = i
            while i < len(line) and line[i] != '\x02': i += 1
            segments.append(('bolditalic', line[start:i]))
            i += 1
        elif line[i] == '\x03':
            i += 1
            start = i
            while i < len(line) and line[i] != '\x04': i += 1
            segments.append(('bold', line[start:i]))
            i += 1
        elif line[i] == '\x05':
            i += 1
            start = i
            while i < len(line) and line[i] != '\x06': i += 1
            segments.append(('italic', line[start:i]))
            i += 1
        else:
            start = i
            while i < len(line) and line[i] not in '\x01\x03\x05': i += 1
            if i > start:
                segments.append(('text', line[start:i]))

    # Restore literal asterisks
    segments = [(style, text.replace('\x07', '*')) for style, text in segments]
    return segments

def font_for_style(style, font, font_bold, font_italic, font_bolditalic):
    if style == 'bolditalic':
        return font_bolditalic or font_bold or font_italic or font
    elif style == 'bold':
        return font_bold
    elif style == 'italic':
        return font_italic
    else:
        return font

def wrap_segments(segments, font, font_bold, font_italic, font_bolditalic, max_width, start_x=0):
    line = []
    x = start_x
    for style, text in segments:
        words = re.split(r'(\s+)', text)
        for word in words:
            if word == '':
                continue
            f = font_for_style(style, font, font_bold, font_italic, font_bolditalic)
            w = f.getbbox(word)[2] - f.getbbox(word)[0]
            if x + w > max_width and line:
                yield line
                line = []
                x = start_x
            line.append((style, word))
            x += w
    if line:
        yield line

def render(md, dithering, printmode, uploaded_img_bytes=None, bannermode=False):
    font = ImageFont.truetype(FONT_PATH, FONT_SIZE)
    font_bold = ImageFont.truetype(FONT_PATH_BOLD, FONT_SIZE)
    font_italic = ImageFont.truetype(FONT_PATH_OBLIQUE, FONT_SIZE)
    font_banner = ImageFont.truetype(FONT_PATH_BOLD, BANNER_FONT_SIZE)
    try:
        font_bolditalic = ImageFont.truetype(FONT_PATH_BOLDITALIC, FONT_SIZE)
    except:
        font_bolditalic = None
    font_h1 = ImageFont.truetype(FONT_PATH_BOLD, HEADER_SIZE_1)
    font_h2 = ImageFont.truetype(FONT_PATH_BOLD, HEADER_SIZE_2)
    font_h3 = ImageFont.truetype(FONT_PATH_BOLD, HEADER_SIZE_3)

    if bannermode:
        # Remove line breaks for single-line banner
        md = md.replace('\r\n', ' ') # should change this if we're ever on windows
        # Only render as a single text line (ignore markdown except bold/italic)
        segments = parse_segments(md)
        # Calculate total width needed
        x = 0
        for style, text in segments:
            f = font_banner
            x += f.getbbox(text)[2] - f.getbbox(text)[0]
        width = max(x, 1)
        # Create image: width x 384 (height is printer width)
        image = Image.new("L", (width, IMAGE_WIDTH), 255)
        draw = ImageDraw.Draw(image)
        x = 0
        y = (IMAGE_WIDTH - BANNER_FONT_SIZE) // 2  # Center vertically
        for style, text in segments:
            f = font_banner
            draw.text((x, y), text, font=f, fill=0)
            x += f.getbbox(text)[2] - f.getbbox(text)[0]
        # Rotate so text is vertical
        image = image.rotate(270, expand=True)
        return image
    
    lines_out = []
    
    for src_line in md.splitlines():
        if src_line.strip() == '':
            lines_out.append(('blank', []))
            continue
        tag = parse_line(src_line)            
        if tag[0] == 'image':
            try:
                image = bleh_image_from_url(tag[1], dithering, printmode)
                lines_out.append(('image', image))
            except Exception as e:
                lines_out.append(('text', [('text', f"[Imagen inválida: {e}]")], font, FONT_SIZE))
        elif tag[0] == 'userimage':
          if uploaded_img_bytes:
            try:
                image = bleh_image_from_bytes(uploaded_img_bytes, dithering, printmode)
                lines_out.append(('image', image))
            except Exception as e:
                lines_out.append(('text', [('text', f"[Error al procesar imagen]")], font, FONT_SIZE))
                print(f"Image processing error: {e}")
          else:
            lines_out.append(('text', [('text', "[No se subió una imagen]")], font, FONT_SIZE))
        elif tag[0] == 'hr':
            lines_out.append(('hr',))
        elif tag[0] == 'header':
            header_level = tag[1]
            segments = tag[2]
            if header_level == 1:
                font_h = font_h1
                size_h = HEADER_SIZE_1
            elif header_level == 2:
                font_h = font_h2
                size_h = HEADER_SIZE_2
            else:
                font_h = font_h3
                size_h = HEADER_SIZE_3
            for wrapped in wrap_segments(segments, font_h, font_h, font_h, font_h, IMAGE_WIDTH):
                lines_out.append(('header', wrapped, font_h, size_h))
        elif tag[0] == 'bullet':
            segments = tag[1]
            bullet_font = font
            bullet_w = bullet_font.getbbox(BULLET_CHAR)[2] - bullet_font.getbbox(BULLET_CHAR)[0]
            wrapped_lines = list(wrap_segments(segments, font, font_bold, font_italic, font_bolditalic, IMAGE_WIDTH - bullet_w, start_x=0))
            for i, wrapped in enumerate(wrapped_lines):
                if i == 0:
                    lines_out.append(('bullet', wrapped, bullet_font, FONT_SIZE, True, bullet_w))
                else:
                    lines_out.append(('bullet', wrapped, bullet_font, FONT_SIZE, False, bullet_w))
        elif tag[0] == 'ordered':
            idx, segments = tag[1], tag[2]
            num_str = f"{idx}. "
            number_font = font
            num_w = number_font.getbbox(num_str)[2] - number_font.getbbox(num_str)[0]
            wrapped_lines = list(wrap_segments(segments, font, font_bold, font_italic, font_bolditalic, IMAGE_WIDTH - num_w, start_x=0))
            for i, wrapped in enumerate(wrapped_lines):
                if i == 0:
                    lines_out.append(('ordered', wrapped, number_font, FONT_SIZE, num_str, True, num_w))
                else:
                    lines_out.append(('ordered', wrapped, number_font, FONT_SIZE, num_str, False, num_w))
        else:  # normal text
            segments = tag[1]
            for wrapped in wrap_segments(segments, font, font_bold, font_italic, font_bolditalic, IMAGE_WIDTH):
                lines_out.append(('text', wrapped, font, FONT_SIZE))

    # Compute total height, including images
    height = 10  # Top margin
    for item in lines_out:
        if item[0] in ('header', 'text', 'bullet', 'ordered'):
            height += item[3]
        elif item[0] == 'hr':
            height += 10
        elif item[0] == 'blank':
            height += FONT_SIZE
        elif item[0] == 'image':
            img = item[1]
            height += img.height + 10  # add margin below image

    image = Image.new("L", (IMAGE_WIDTH, height), 255)
    draw = ImageDraw.Draw(image)
    y = 0
    for item in lines_out:
        if item[0] == 'blank':
            y += FONT_SIZE
        elif item[0] == 'hr':
            draw.line((0, y + 5, IMAGE_WIDTH, y + 5), fill=0, width=2)
            y += 10
        elif item[0] == 'header':
            segments, fnt, sz = item[1], item[2], item[3]
            x = 0
            for style, text in segments:
                draw.text((x, y), text, font=fnt, fill=0)
                x += fnt.getbbox(text)[2] - fnt.getbbox(text)[0]
            y += sz
        elif item[0] == 'bullet':
            segments, bullet_font, sz, show_bullet, bullet_w = item[1], item[2], item[3], item[4], item[5]
            x = 0
            if show_bullet:
                draw.text((x, y), BULLET_CHAR, font=bullet_font, fill=0)
                x += bullet_w
            else:
                x += bullet_w
            for style, text in segments:
                f = font_for_style(style, bullet_font, font_bold, font_italic, font_bolditalic)
                draw.text((x, y), text, font=f, fill=0)
                x += f.getbbox(text)[2] - f.getbbox(text)[0]
            y += sz
        elif item[0] == 'ordered':
            segments, number_font, sz, num_str, show_num, num_w = item[1], item[2], item[3], item[4], item[5], item[6]
            x = 0
            if show_num:
                draw.text((x, y), num_str, font=number_font, fill=0)
                x += num_w
            else:
                x += num_w
            for style, text in segments:
                f = font_for_style(style, number_font, font_bold, font_italic, font_bolditalic)
                draw.text((x, y), text, font=f, fill=0)
                x += f.getbbox(text)[2] - f.getbbox(text)[0]
            y += sz
        elif item[0] == 'text':
            segments, fnt, sz = item[1], item[2], item[3]
            x = 0
            for style, text in segments:
                f = font_for_style(style, font, font_bold, font_italic, font_bolditalic)
                draw.text((x, y), text, font=f, fill=0)
                x += f.getbbox(text)[2] - f.getbbox(text)[0]
            y += sz
        elif item[0] == 'image':
            img = item[1]
            # Center image horizontally if narrower
            img_x = (IMAGE_WIDTH - img.width) // 2 if img.width < IMAGE_WIDTH else 0
            image.paste(img, (img_x, y))
            y += img.height + 10  # vertical margin after image
    if bannermode:
        image = image.rotate(270, expand=True)
    return image

@app.route("/", methods=["GET", "POST"])
def index():
    img_data = None
    md = ""
    printed = False
    error = None
    lang = request.args.get('lang') or session.get('lang', 'en')
    session['lang'] = lang
    t = TRANSLATIONS.get(lang, TRANSLATIONS['en'])
    dithering_modes = {
      t['no_dithering']: "none",
      "Floyd-Steinberg": "floyd",
      "Bayer 2x2": "bayer2x2",
      "Bayer 4x4": "bayer4x4",
      "Bayer 8x8": "bayer8x8",
      "Bayer 16x16": "bayer16x16",
      "Atkinson": "atkinson",
      "Jarvis-Judice-Ninke": "jjn"
    }

    if request.method == "POST":
        userimg = request.files.get("userimg")
        uploaded_img_bytes = None
        rotation = int(request.form.get("rotation", "0"))
        if userimg and userimg.filename:
            # Remove old temp file if present
            remove_uploaded_img()
            # Save new file to a temp location
            with tempfile.NamedTemporaryFile(delete=False, suffix=".catnote") as tmpf:
                userimg.save(tmpf)
                tmpf.flush()
                session['uploaded_img_path'] = tmpf.name
                tmpf.seek(0)
                uploaded_img_bytes = tmpf.read()
        elif 'uploaded_img_path' in session:
            path = session['uploaded_img_path']
            if os.path.exists(path):
                with open(path, "rb") as f:
                    uploaded_img_bytes = f.read()
            else:
                # If file was deleted or missing, remove from session
                session.pop('uploaded_img_path', None)
        else:
            remove_uploaded_img()
        if uploaded_img_bytes and rotation in (90, 180, 270):
            try:
                uploaded_img_bytes = rotate_image_bytes(uploaded_img_bytes, rotation)
            except Exception as e:
                error = f"Error rotating image: {e}"
                uploaded_img_bytes = None
        md = request.form["md"]
        dithering = request.form.get("dithering", "floyd")
        printmode = request.form.get("printmode", "1bpp")
        bannermode = bool(request.form.get("bannermode"))
        image = render(md, dithering, printmode, uploaded_img_bytes, bannermode=bannermode)
        session['dithering'] = dithering
        session['printmode'] = printmode
        session['rotation'] = rotation
        session['bannermode'] = bannermode
        buf = io.BytesIO()
        image.save(buf, format="PNG")
        buf.seek(0)
        img_data = base64.b64encode(buf.getvalue()).decode()
        intensity = request.form.get("intensity", "85")
        try:
            intensity = int(intensity)
            if intensity < 0 or intensity > 100:
                raise ValueError("Intensity must be between 0 and 100")
        except ValueError:
            error = "Intensity must be a number between 0 and 100"
            intensity = 85 
        session['intensity'] = intensity
        # If print button pressed, send to blehd directly via Unix socket
        if "print" in request.form:
            try:
                printable = pad_image_to_min_lines(image)
                pixels, height = pack_pixels(printable, printmode)
                params = {
                    "mode": printmode,
                    "intensity": intensity,
                    "height": height,
                    "pixels_b64": base64.b64encode(pixels).decode("ascii"),
                }
                blehd_call("print.start", params)
                printed = True
            except Exception as e:
                error = f"Failed to print via blehd socket: {e}"
    return render_template(
        "index.html",
        dithering_modes=dithering_modes,
        img=img_data,
        default_md=md,
        printed=printed,
        error=error,
        current_intensity=session.get('intensity', 85),
        current_dithering=session.get('dithering', 'floyd'),
        current_rotation=session.get('rotation', 0),
        current_printmode=session.get('printmode', '1bpp'),
        current_bannermode=session.get('bannermode', False),
        t=t,
        lang=lang
    )

@app.route('/manifest.json')
def manifest():
    return send_from_directory('static', 'manifest.json')

@app.route('/icon512_maskable.png')
def icon_maskable():
    return send_from_directory('static', 'icon512_maskable.png')

@app.route('/icon512_rounded.png')
def icon_rounded():
    return send_from_directory('static', 'icon512_rounded.png')

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='CatNote: Server for Markdown to MXW01 Cat Printer')
    parser.add_argument('-p', '--port', type=int, default=5000, help='Port to run the server on (default: 5000)')
    args = parser.parse_args()
    app.run(host='0.0.0.0', port=args.port)