from __future__ import annotations

# :asciigram: disable
import argparse
import dataclasses
import hashlib
import json
import re
import subprocess
import sys
from pathlib import Path
from typing import TypedDict


class AsciigramDisabled(Exception):
    pass


@dataclasses.dataclass
class Context:
    file_ext: str
    interactive: bool = False
    block_marker: str = ":asciigram:"


class Options(TypedDict, total=False):
    lang: str | None  # May not be needed if command is set
    command: list[str] | None  # Might be drawn from lang
    code_start: str | None
    code_end: str | None
    ascii_start: str | None
    ascii_end: str | None
    ascii_header: list[str] | None
    ascii_footer: list[str] | None
    prefix: str | None
    # internal
    _hash: str | None


class Config(TypedDict):
    defaults: Options
    overrides_for_file_ext: dict[str, Options]
    command_for_lang: dict[str, list[str]]


BASE_CONFIG: Config = Config(
    defaults=Options(
        code_start="```(?P<lang>(?!(?:text|ascii)$).+)$",
        code_end="```$",
        ascii_start=None,
        ascii_end=None,
        ascii_header=None,
        ascii_footer=None,
        prefix=None,
    ),
    overrides_for_file_ext={
        ".md": Options(
            ascii_start="```text$",
            ascii_end="```$",
            ascii_header=["```text"],
            # The multiple-footer thing doesn't actually work well -- very hard to
            # properly splice out the footer because it's not part of the ascii_end regex
            # Let's get rid of it for now.
            ascii_footer=["```"],
        )
    },
    command_for_lang={
        "d2": ["d2", "-", "--stdout-format", "txt", "-"],
        "plantuml": ["plantuml", "-pipe", "-utxt"],
    },
)


def extract_content(
    lines: list[str], start_index: int, end_index: int, prefix: str | None = None
) -> str:
    """Extract the context removing the prefix."""
    if prefix is not None:
        cleaned_lines = []
        for line in lines[start_index:end_index]:
            cleaned_lines.append(line.removeprefix(prefix))
    else:
        cleaned_lines = lines[start_index:end_index]
    return "\n".join(cleaned_lines)


def generate_ascii(command: list[str], input: str) -> list[str]:
    rv: list[str]
    try:
        result = subprocess.run(
            command, input=input, capture_output=True, text=True, check=True
        )
        rv = result.stdout
    except subprocess.CalledProcessError as e:
        rv = f"Error running command {command}: {e.stderr}"
    except FileNotFoundError as e:
        rv = f"Command {command} not found: {e}"
    except Exception as e:
        rv = f"Error running command {command}: {e}"
    return rv.splitlines()


def get_prefix(line: str, regex: str) -> str:
    match = re.search(regex, line)
    assert match is not None, f"No match found for regex {regex} in line {line}"
    start, _ = match.span()
    return line[:start]


def wrap(lines: list[str], prefix: str, postfix: str) -> list[str]:
    result = []
    for line in lines:
        result.append(f"{prefix}{line}{postfix}")
    return result


def remove_prefix(prefix: str, lines: list[str]) -> list[str]:
    result = []
    for line in lines:
        result.append(line.removeprefix(prefix))
    return result


def find_next_matching_index(
    lines: list[str], regex: str, start_index: int
) -> int | None:
    for i in range(start_index, len(lines)):
        if re.search(regex, lines[i]):
            return i
    return None


def get_first_block_by_start_end_delimiters(
    lines: list[str], start_regex: str, end_regex: str, start_index: int = 0
) -> tuple[int, int] | None:
    start_index = find_next_matching_index(lines, start_regex, start_index)
    if start_index is None:
        return None
    end_index = find_next_matching_index(lines, end_regex, start_index)
    if end_index is None:
        return None
    return (start_index, end_index)


def get_blocks_by_single_delimiter(
    lines: list[str], block_marker_regex: str
) -> list[tuple[int, int]]:
    blocks: list[tuple[int, int]] = []
    start_index = find_next_matching_index(lines, block_marker_regex, 0)
    while start_index is not None:
        end_index = find_next_matching_index(lines, block_marker_regex, start_index + 1)
        if end_index is None:
            blocks.append((start_index, len(lines)))
            break
        else:
            blocks.append((start_index, end_index))
            start_index = end_index
    return blocks


def get_lang(line: str, regex: str) -> str | None:
    match = re.search(regex, line)
    if match is None:
        return None
    return match.group("lang")


def load_json_inline(marker_regex: str, line: str) -> Options | None:
    match = re.search(marker_regex, line)
    if match is None:
        return None
    _, end = match.span()
    options = line[end:].strip()
    if options == "disable":
        raise AsciigramDisabled
    options_start = line.find("{", end)
    existing_options = {}
    if options_start >= 0:
        existing_options, _ = json.JSONDecoder().raw_decode(line, options_start)
    return Options(**existing_options)


def generate_hash(content: str) -> str:
    return hashlib.sha256(content.encode()).hexdigest()[:8]


def save_json_inline(marker_regex: str, options: dict, line: str) -> list[str]:
    marker = re.search(marker_regex, line)
    assert marker is not None, f"No block marker found in line {line}"
    _, marker_end = marker.span()
    options_start = line.find("{", marker_end)
    if options_start < 0:
        options_start = options_end = marker_end
        pad = " "
    else:
        _, options_end = json.JSONDecoder().raw_decode(line, options_start)
        pad = ""

    options_str = json.dumps(options, sort_keys=True)
    return line[:options_start] + pad + options_str + line[options_end:]


def merge_options(*options: Options) -> Options:
    """Apply overrides from left to right. Skips None values."""
    result = {}
    for option in options:
        for key, value in option.items():
            if value is not None:
                result[key] = value
    return Options(**result)


def transform_one(
    context: Context, default_options: Options, lines: list[str]
) -> list[str]:
    line_options = load_json_inline(context.block_marker, lines[0])
    if line_options is None:
        raise ValueError(f"Malformed block marker: {lines[0]}")

    # We need to find the code block to get the rest of the options.
    initial_options = merge_options(default_options, line_options)
    code_start = initial_options.get("code_start", None)
    code_end = initial_options.get("code_end", None)
    if code_start is None or code_end is None:
        raise ValueError(f"Code start or end is required: {initial_options}")

    code_block = get_first_block_by_start_end_delimiters(
        lines, code_start, code_end, start_index=1
    )

    if code_block is None:
        raise ValueError(f"No code block found for {initial_options}")

    first_line_of_code_block = lines[code_block[0]]
    code_lang = get_lang(first_line_of_code_block, code_start)
    code_prefix = get_prefix(first_line_of_code_block, code_start)
    code_options = Options(lang=code_lang, prefix=code_prefix)

    final_options = merge_options(default_options, code_options, line_options)

    print(f"final_options = {final_options}")

    ascii_command = final_options.get("command", None)
    if ascii_command is None:
        # TODO: Think of a way to avoid the global lookup -- maybe move to context?
        ascii_command = BASE_CONFIG["command_for_lang"].get(code_lang, None)
    if ascii_command is None:
        raise ValueError(
            f"No ASCII command found! default: {default_options}, code: {code_options}, line: {line_options}"
        )

    prefix = final_options.get("prefix", None)
    if prefix is None:
        raise ValueError(f"Prefix is required: {final_options}")

    # Do not include delimiters in the code.
    code = extract_content(lines, code_block[0] + 1, code_block[1], prefix)

    code_hash = generate_hash(code)
    existing_hash = final_options.get("_hash", None)
    if context.interactive:
        print(f"Existing hash: {existing_hash}, new hash: {code_hash}")

    if existing_hash is not None and existing_hash == code_hash:
        return lines

    ascii_block = None
    ascii_start = final_options.get("ascii_start", None)
    ascii_end = final_options.get("ascii_end", None)

    if ascii_start is None:
        print("ascii_start is None, using code block")
        assert ascii_end is None, "ascii_end must be None if ascii_start is None"
        ascii_block = (1, code_block[0] - 1)
    else:
        print("ascii_start is not None, using ascii block")
        assert ascii_end is not None, "ascii_end is required if ascii_start is set"
        ascii_block = get_first_block_by_start_end_delimiters(
            lines, ascii_start, ascii_end, start_index=1
        )
        # If there's no ascii block yet, but we expect to have one, then insert directly
        # under the first line.
        if ascii_block is None:
            print("no ascii block found, inserting directly under the first line")
            ascii_block = (0, 0)
        else:
            print(f"found ascii block = {ascii_block}")

    print(f"found ascii_block = {ascii_block}")

    if context.interactive:
        print(f"Running command {ascii_command} on code block")
    ascii = generate_ascii(ascii_command, code)

    ascii_header = final_options.get("ascii_header") or []
    ascii_footer = final_options.get("ascii_footer") or []

    ascii = ascii_header + ascii + ascii_footer

    line_options["_hash"] = code_hash
    new_marker = [save_json_inline(context.block_marker, line_options, lines[0])]

    return (
        new_marker
        + lines[1 : ascii_block[0]]
        + wrap(ascii, prefix=prefix, postfix="\n")
        + lines[ascii_block[1] + 1 :]
    )


def validate_options(options: Options) -> list[str]:
    errors = []
    if options.get("code_start", None) is None:
        errors.append("code_start is required")
    if options.get("code_end", None) is None:
        errors.append("code_end is required")
    return errors


def transform_all(context: Context, user_options: Options, text: str) -> str:
    default_options = BASE_CONFIG["defaults"].copy()
    file_overrides = BASE_CONFIG["overrides_for_file_ext"].get(context.file_ext, {})
    print(f"file_overrides = {file_overrides}")
    default_options = merge_options(default_options, file_overrides, user_options)

    original_lines = text.splitlines(keepends=True)
    result = []
    prev_end_index = 0
    blocks = get_blocks_by_single_delimiter(original_lines, context.block_marker)
    # check for disable marker
    for start_index, _ in blocks:
        # hack: will raise AsciigramDisabled if disable marker is found
        load_json_inline(context.block_marker, original_lines[start_index])
    for start_index, end_index in blocks:
        result += original_lines[prev_end_index:start_index]
        result += transform_one(
            context, default_options, original_lines[start_index:end_index]
        )
        prev_end_index = end_index
    result += original_lines[prev_end_index:]
    return "".join(result)


def main():
    parser = argparse.ArgumentParser(description="Transform code blocks to ASCII art")
    parser.add_argument(
        "file",
        type=Path,
        help="Input file to process",
    )
    parser.add_argument(
        "-o",
        "--output-file",
        type=Path,
        default=None,
        help="Output file (default: overwrite input file)",
    )
    args = parser.parse_args()

    input_file = args.file
    if input_file.is_dir():
        sys.exit(0)

    context = Context(file_ext=input_file.suffix)

    print(f"context: {context}")

    text = input_file.read_text()

    # TODO: Make it more efficient. What do other formatters do?
    try:
        out = transform_all(context, {}, text)
    except AsciigramDisabled:
        print(":asciigram: disable found, skipping transformation")
        sys.exit(0)

    if out == text:
        # Very important. Or else watchexec and the rest will spin forever.
        sys.exit(0)

    output_file = args.output_file if args.output_file is not None else input_file

    # TODO: Preserve permissions, metadata, etc. What do other formatters do?
    with open(output_file, "w") as f:
        f.write(out)


if __name__ == "__main__":
    main()
