import argparse
import json

import colorspacious
import numpy as np

# Original values for the Solarized color scheme,
# created by Ethan Schoonover (https://ethanschoonover.com/solarized/)
SOLARIZED_LAB = np.array(
    [
        [15, -12, -12],
        [20, -12, -12],
        [45, -7, -7],
        [50, -7, -7],
        [60, -6, -3],
        [65, -5, -2],
        [92, -0, 10],
        [97, 0, 10],
        [50, 65, 45],
        [50, 50, 55],
        [60, 10, 65],
        [60, -20, 65],
        [60, -35, -5],
        [55, -10, -45],
        [50, 15, -45],
        [50, 65, -5],
    ]
)

# I couldn't get a perfect translation of Solarized L*a*b values into sRGB,
# so here is upstream's translation for reference
SOLARIZED_RGB = np.array(
    [
        [0, 43, 54],
        [7, 54, 66],
        [88, 110, 117],
        [101, 123, 131],
        [131, 148, 150],
        [147, 161, 161],
        [238, 232, 213],
        [253, 246, 227],
        [220, 50, 47],
        [203, 75, 22],
        [181, 137, 0],
        [133, 153, 0],
        [42, 161, 152],
        [38, 139, 210],
        [108, 113, 196],
        [211, 54, 130],
    ]
)

# Parse arguments
parser = argparse.ArgumentParser(
    description="Generate a base16-theme based derived from Solarized"
)
parser.add_argument("--source", choices=["lab", "rgb"], default="lab")
parser.add_argument("--lightness_factor", type=float, default=1.0)
parser.add_argument("--chroma-factor", type=float, default=1.0)
parser.add_argument("--hue_shift", type=float, default=-75.0)
parser.add_argument("--polarity", choices=["dark", "light"], default="dark")
parser.add_argument(
    "--output", choices=["json", "truecolor"], default="truecolor"
)
args = parser.parse_args()

# Convert source to JCh color space
if args.source == "lab":
    solarized_jch = colorspacious.cspace_convert(
        SOLARIZED_LAB, "CIELab", "JCh"
    )
elif args.source == "rgb":
    solarized_jch = colorspacious.cspace_convert(
        SOLARIZED_RGB, "sRGB255", "JCh"
    )

# Build frogarized theme
jch_factor = [args.lightness_factor, args.chroma_factor, 1]
jch_shift = [0, 0, args.hue_shift]
frogarzied_jch = np.vstack(
    [solarized_jch[:8] * jch_factor + jch_shift, solarized_jch[8:]]
)

# Convert frogarized to RGB
frogarized_srgb = colorspacious.cspace_convert(
    frogarzied_jch, "JCh", "sRGB255"
)
frogarized_rgb = np.uint8(np.rint(np.clip(frogarized_srgb, 0, 255)))
if args.polarity == "light":
    frogarized_rgb = np.vstack([frogarized_rgb[7::-1], frogarized_rgb[8:]])

# Output
palette = dict()
for i in range(16):
    rgb = frogarized_rgb[i]
    r, g, b = rgb
    hex = f"#{r:02x}{g:02x}{b:02x}"
    palette[f"base{i:02X}"] = hex
    if args.output == "truecolor":
        print(f"\033[48;2;{r};{g};{b}m{hex}\033[0m")  # ]]
        # treesitter is silly and will consider brackets in strings
        # as indentation, hence the comment above
if args.output == "json":
    scheme = palette.copy()
    scheme.update(
        {
            "slug": f"frogarized-{args.polarity}",
            "scheme": f"Frogarized {args.polarity.title()}",
            "author": "Geoffrey Frogeye (with work from Ethan Schoonover)",
        }
    )
    print(json.dumps(scheme, indent=4))