from __future__ import annotations

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

from watchdog.events import FileSystemEventHandler
from watchdog.observers import Observer

logger = logging.getLogger(__name__)


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)

    logger.info(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:
        logger.info(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:
        logger.info("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:
        logger.info("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:
            logger.info("no ascii block found, inserting directly under the first line")
            ascii_block = (0, 0)
        else:
            logger.info(f"found ascii block = {ascii_block}")

    logger.info(f"found ascii_block = {ascii_block}")

    if context.interactive:
        logger.info(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, {})
    logger.info(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)


class AsciigramHandler(FileSystemEventHandler):
    """Watchdog handler for asciigram file processing."""
    
    def __init__(self, watch_path: Path, output_file: Path | None = None):
        self.watch_path = watch_path.resolve()
        self.output_file = output_file
        self.is_file = watch_path.is_file()
    
    def on_modified(self, event):
        if event.is_directory:
            return
        
        event_path = Path(event.src_path)
        
        # Skip hidden files
        if any(part.startswith(".") for part in event_path.parts):
            return
        
        if self.is_file:
            # Watching a single file - only process if it's the watched file
            if event_path.resolve() == self.watch_path.resolve():
                logger.info(f"Processing {event_path}")
                process_file(event_path, self.output_file)
        else:
            # Watching a directory - process all non-hidden files
            logger.info(f"Change detected, processing all files in {self.watch_path}")
            for file_path in self.watch_path.rglob("*"):
                if file_path.is_file() and not any(part.startswith(".") for part in file_path.parts):
                    process_file(file_path)
    
    def on_created(self, event):
        # Treat created files the same as modified
        self.on_modified(event)


def process_file(input_file: Path, output_file: Path | None = None) -> bool:
    """Process a single file. Returns True if file was modified, False otherwise."""
    if input_file.is_dir():
        return False

    context = Context(file_ext=input_file.suffix)

    text = input_file.read_text()

    try:
        out = transform_all(context, {}, text)
    except AsciigramDisabled:
        logger.info(f":asciigram: disable found, skipping {input_file}")
        return False

    if out == text:
        return False

    output = output_file if output_file is not None else input_file

    with open(output, "w") as f:
        f.write(out)
    
    return True


def main():
    parser = argparse.ArgumentParser(description="Transform code blocks to ASCII art")
    parser.add_argument(
        "file",
        type=Path,
        help="Input file or directory to process",
    )
    parser.add_argument(
        "-o",
        "--output-file",
        type=Path,
        default=None,
        help="Output file (default: overwrite input file)",
    )
    parser.add_argument(
        "-w", "--watch",
        action="store_true",
        help="Watch for file changes and automatically reprocess",
    )
    parser.add_argument(
        "-v",
        action="count",
        default=0,
        help="Increase verbosity (-v for INFO, -vv for DEBUG)",
    )
    args = parser.parse_args()
    
    # Set log level based on verbosity: 0=WARNING, 1=INFO, 2+=DEBUG
    if args.v == 0:
        log_level = logging.WARNING
    elif args.v == 1:
        log_level = logging.INFO
    else:
        log_level = logging.DEBUG
    
    logging.basicConfig(
        level=log_level,
        format="%(message)s",
        stream=sys.stderr,
    )

    input_path = args.file

    if args.watch:
        # Initial processing
        if input_path.is_file():
            logger.info(f"Watching {input_path} for changes...")
            process_file(input_path, args.output_file)
        else:
            logger.info(f"Watching {input_path} recursively for changes...")
            for file_path in input_path.rglob("*"):
                if file_path.is_file() and not any(part.startswith(".") for part in file_path.parts):
                    process_file(file_path)
        
        # Set up watch
        handler = AsciigramHandler(input_path, args.output_file)
        observer = Observer()
        
        if input_path.is_file():
            observer.schedule(handler, str(input_path.parent), recursive=False)
        else:
            observer.schedule(handler, str(input_path), recursive=True)
        
        observer.start()
        try:
            observer.join()
        except KeyboardInterrupt:
            observer.stop()
        observer.join()
        return
    
    # Non-watch mode: process the file or directory
    if input_path.is_dir():
        # Process all non-hidden files in directory
        processed = 0
        for file_path in input_path.rglob("*"):
            if file_path.is_file() and not any(part.startswith(".") for part in file_path.parts):
                if process_file(file_path):
                    processed += 1
        sys.exit(0)
    else:
        # Process single file
        process_file(input_path, args.output_file)


if __name__ == "__main__":
    main()
