#!/usr/bin/env python
"""
Creation of templates for Exercitium problems
"""

import ast
import re
from argparse import ArgumentParser
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path


@dataclass(frozen=True)
class Example:
    input_values: tuple[object, ...]
    expected_value: object


def load_instructions(path: Path) -> str:
    return path.read_text(encoding="utf-8").strip()

def split_top_level_values(text: str) -> list[str]:
    parts: list[str] = []
    current: list[str] = []
    depth = 0
    quote_char: str | None = None
    escaped = False

    for char in text.strip():
        if quote_char is not None:
            current.append(char)
            if escaped:
                escaped = False
            elif char == "\\":
                escaped = True
            elif char == quote_char:
                quote_char = None
            continue

        if char in {"'", '"'}:
            quote_char = char
            current.append(char)
            continue

        if char in "([{":
            depth += 1
            current.append(char)
            continue

        if char in ")]}":
            depth -= 1
            current.append(char)
            continue

        if char.isspace() and depth == 0:
            if current:
                parts.append("".join(current))
                current = []
            continue

        current.append(char)

    if quote_char is not None or depth != 0:
        raise ValueError(f"Cannot parse input values: {text}")

    if current:
        parts.append("".join(current))

    return parts


def parse_input_values(text: str) -> tuple[object, ...]:
    parts = split_top_level_values(text)
    if not parts:
        raise ValueError("No input values found")

    return tuple(ast.literal_eval(part) for part in parts)

def parse_instructions(text: str) -> tuple[str, list[Example]]:
    examples: list[Example] = []
    function_name: str | None = None

    for raw_line in text.splitlines():
        line = raw_line.strip()
        if not line:
            continue

        left, right = line.split("==", maxsplit=1)
        match = re.match(r"^--\s*([A-Za-z_][A-Za-z0-9_]*)\s+(.*)$", left.strip())
        if match is None:
            raise ValueError(f"Cannot parse instruction line: {raw_line}")

        current_name = match.group(1)
        if function_name is None:
            function_name = current_name
        elif function_name != current_name:
            raise ValueError("All instruction lines must use the same function name")

        examples.append(
            Example(
                input_values=parse_input_values(match.group(2).strip()),
                expected_value=ast.literal_eval(right.strip()),
            )
        )

    if function_name is None:
        raise ValueError("No examples found in instructions")

    return function_name, examples

def build_names(prefix: str, count: int) -> list[str]:
    return [f"{prefix}_{index}" for index in range(1, count + 1)]

def infer_python_type(value: object) -> str:
    if isinstance(value, bool):
        return "bool"
    if isinstance(value, int):
        return "int"
    if isinstance(value, str):
        return "str"
    if isinstance(value, list):
        inner_type = infer_python_type(value[0]) if value else "object"
        return f"list[{inner_type}]"
    if isinstance(value, tuple):
        inner_types = ", ".join(infer_python_type(item) for item in value)
        return f"tuple[{inner_types}]"
    return "object"


def infer_go_type(value: object) -> str:
    if isinstance(value, bool):
        return "bool"
    if isinstance(value, int):
        return "int"
    if isinstance(value, str):
        return "string"
    if isinstance(value, list):
        inner_type = infer_go_type(value[0]) if value else "int"
        return f"[]{inner_type}"
    raise TypeError(f"Unsupported Go type: {value!r}")


def infer_rust_type(value: object) -> str:
    if isinstance(value, bool):
        return "bool"
    if isinstance(value, int):
        return "i32"
    if isinstance(value, str):
        return "&str"
    if isinstance(value, list):
        inner_type = infer_rust_type(value[0]) if value else "i32"
        return f"Vec<{inner_type}>"
    if isinstance(value, tuple):
        inner_types = ", ".join(infer_rust_type(item) for item in value)
        return f"({inner_types})"
    raise TypeError(f"Unsupported Rust type: {value!r}")


def infer_rust_param_type(value: object) -> str:
    if isinstance(value, list):
        inner_type = infer_rust_type(value[0]) if value else "i32"
        return f"&[{inner_type}]"
    return infer_rust_type(value)


def render_python(value: object) -> str:
    return repr(value)


def render_julia(value: object) -> str:
    if isinstance(value, str):
        return f'"{value}"'
    if isinstance(value, list):
        return "[" + ", ".join(render_julia(item) for item in value) + "]"
    if isinstance(value, tuple):
        return "(" + ", ".join(render_julia(item) for item in value) + ")"
    return str(value)


def render_go(value: object) -> str:
    if isinstance(value, str):
        return f'"{value}"'
    if isinstance(value, list):
        return (
            f"{infer_go_type(value)}{{"
            + ", ".join(render_go(item) for item in value)
            + "}"
        )
    return str(value)


def render_rust(value: object) -> str:
    if isinstance(value, str):
        return f'"{value}"'
    if isinstance(value, list):
        return "vec![" + ", ".join(render_rust(item) for item in value) + "]"
    if isinstance(value, tuple):
        return "(" + ", ".join(render_rust(item) for item in value) + ")"
    return str(value)


def build_python_checks(function_name: str, examples: list[Example]) -> str:
    blocks: list[str] = []

    for example in examples:
        lines: list[str] = []
        names = build_names("check", len(example.input_values))

        for name, value in zip(names, example.input_values):
            lines.append(f"{name} = {render_python(value)}")

        lines.append(
            f"print({function_name}({', '.join(names)}))"
            f"  # {render_python(example.expected_value)}"
        )
        blocks.append("\n".join(lines))

    return "\n\n".join(blocks)


def build_julia_checks(function_name: str, examples: list[Example]) -> str:
    blocks: list[str] = []

    for example in examples:
        lines: list[str] = []
        names = build_names("check", len(example.input_values))

        for name, value in zip(names, example.input_values):
            lines.append(f"{name} = {render_julia(value)}")

        lines.append(
            f"println({function_name}({', '.join(names)}))"
            f"  # {render_julia(example.expected_value)}"
        )
        blocks.append("\n".join(lines))

    return "\n\n".join(blocks)

def render_rust_call_arg(name: str, value: object) -> str:
    if isinstance(value, list):
        return f"&{name}"
    return name

def build_rust_checks(function_name: str, examples: list[Example]) -> str:
    blocks: list[str] = []

    for example in examples:
        lines: list[str] = []
        names = build_names("check", len(example.input_values))

        for name, value in zip(names, example.input_values):
            lines.append(f"    let {name} = {render_rust(value)};")

        call_args = ", ".join(
            render_rust_call_arg(name, value)
            for name, value in zip(names, example.input_values)
        )
        lines.append(
            f'    println!("{{:?}}", {function_name}({call_args}));'
            f"  // {render_python(example.expected_value)}"
        )
        blocks.append("\n".join(lines))

    return "\n\n".join(blocks)


def build_go_checks(function_name: str, examples: list[Example]) -> str:
    blocks: list[str] = []

    for example in examples:
        lines: list[str] = ["    {"]
        input_names = build_names("check", len(example.input_values))

        for name, value in zip(input_names, example.input_values):
            lines.append(f"        {name} := {render_go(value)}")

        if isinstance(example.expected_value, tuple):
            result_names = build_names("result", len(example.expected_value))
            lines.append(
                f"        {', '.join(result_names)} := "
                f"{function_name}({', '.join(input_names)})"
            )
            lines.append(
                f"        fmt.Println({', '.join(result_names)})"
                f"  // {render_python(example.expected_value)}"
            )
        else:
            lines.append(
                f"        fmt.Println({function_name}({', '.join(input_names)}))"
                f"  // {render_python(example.expected_value)}"
            )

        lines.append("    }")
        blocks.append("\n".join(lines))

    return "\n\n".join(blocks)

def default_python_value(value: object) -> str:
    if isinstance(value, bool):
        return "False"
    if isinstance(value, int):
        return "0"
    if isinstance(value, str):
        return '""'
    if isinstance(value, list):
        return "[]"
    if isinstance(value, tuple):
        return "(" + ", ".join(default_python_value(item) for item in value) + ")"
    return "None"


def default_julia_value(value: object) -> str:
    if isinstance(value, bool):
        return "false"
    if isinstance(value, int):
        return "0"
    if isinstance(value, str):
        return '""'
    if isinstance(value, list):
        return "[]"
    if isinstance(value, tuple):
        return "(" + ", ".join(default_julia_value(item) for item in value) + ")"
    return "nothing"


def default_rust_value(value: object) -> str:
    if isinstance(value, bool):
        return "false"
    if isinstance(value, int):
        return "0"
    if isinstance(value, str):
        return '""'
    if isinstance(value, list):
        return "vec![]"
    if isinstance(value, tuple):
        return "(" + ", ".join(default_rust_value(item) for item in value) + ")"
    raise TypeError(f"Unsupported Rust type: {value!r}")


def default_go_value(value: object) -> str:
    if isinstance(value, bool):
        return "false"
    if isinstance(value, int):
        return "0"
    if isinstance(value, str):
        return '""'
    if isinstance(value, list):
        return f"{infer_go_type(value)}{{}}"
    if isinstance(value, tuple):
        return ", ".join(default_go_value(item) for item in value)
    raise TypeError(f"Unsupported Go type: {value!r}")


def infer_go_signature_return(value: object) -> str:
    if isinstance(value, tuple):
        return "(" + ", ".join(infer_go_type(item) for item in value) + ")"
    return infer_go_type(value)

def comment_block(prefix: str, text: str) -> str:
    return "\n".join(
        f"{prefix} {line}".rstrip()
        for line in text.splitlines()
    )

def create_if_missing(
    output_path: Path,
    creator: Callable[[Path, str, list[Example]], None],
    function_name: str,
    examples: list[Example],
) -> None:
    if output_path.exists():
        return

    output_path.parent.mkdir(parents=True, exist_ok=True)
    creator(output_path, function_name, examples)

def create_python(
    problem_py: Path,
    function_name: str,
    examples: list[Example],
) -> None:
    input_names = build_names("value", len(examples[0].input_values))
    params = ", ".join(
        f"{name}: {infer_python_type(value)}"
        for name, value in zip(input_names, examples[0].input_values)
    )
    return_type = infer_python_type(examples[0].expected_value)

    template = "\n".join(
        [
            f"def {function_name}({params}) -> {return_type}:",
            f"    return {default_python_value(examples[0].expected_value)}",
            "",
            build_python_checks(function_name, examples),
            "",
        ]
    )
    problem_py.write_text(template, encoding="utf-8")


def create_julia(
    problem_jl: Path,
    function_name: str,
    examples: list[Example],
) -> None:
    params = ", ".join(build_names("value", len(examples[0].input_values)))

    template = "\n".join(
        [
            f"function {function_name}({params})",
            f"    return {default_julia_value(examples[0].expected_value)}",
            "end",
            "",
            build_julia_checks(function_name, examples),
            "",
        ]
    )
    problem_jl.write_text(template, encoding="utf-8")


def create_rust(
    problem_rs: Path,
    function_name: str,
    examples: list[Example],
) -> None:
    input_names = build_names("value", len(examples[0].input_values))
    params = ", ".join(
        f"{name}: {infer_rust_param_type(value)}"
        for name, value in zip(input_names, examples[0].input_values)
    )
    return_type = infer_rust_type(examples[0].expected_value)

    template = "\n".join(
        [
            f"fn {function_name}({params}) -> {return_type} {{",
            f"    {default_rust_value(examples[0].expected_value)}",
            "}",
            "",
            "fn main() {",
            build_rust_checks(function_name, examples),
            "}",
            "",
        ]
    )
    problem_rs.write_text(template, encoding="utf-8")


def create_go(
    problem_go: Path,
    function_name: str,
    examples: list[Example],
) -> None:
    input_names = build_names("value", len(examples[0].input_values))
    params = ", ".join(
        f"{name} {infer_go_type(value)}"
        for name, value in zip(input_names, examples[0].input_values)
    )
    return_type = infer_go_signature_return(examples[0].expected_value)

    template = "\n".join(
        [
            "package main",
            "",
            'import "fmt"',
            "",
            f"func {function_name}({params}) {return_type} " + "{",
            f"    return {default_go_value(examples[0].expected_value)}",
            "}",
            "",
            "func main() {",
            build_go_checks(function_name, examples),
            "}",
            "",
        ]
    )
    problem_go.write_text(template, encoding="utf-8")


if __name__ == "__main__":
    parser = ArgumentParser(description=__doc__)
    parser.add_argument(
        "-p",
        "--problem-number",
        dest="problem_number",
        type=int,
        required=True,
        help="number of the problem to solve",
    )
    args = parser.parse_args()

    base_dir = Path(__file__).resolve().parent
    instructions = load_instructions(base_dir / "instructions.txt")

    problem_stem = f"{args.problem_number:03d}"
    problem_py = base_dir / "Python" / f"{problem_stem}.py"
    problem_jl = base_dir / "Julia" / f"{problem_stem}.jl"
    problem_rs = base_dir / "Rust" / f"{problem_stem}.rs"
    problem_go = base_dir / "Go" / f"{problem_stem}.go"
    
    function_name, examples = parse_instructions(instructions)

    create_if_missing(problem_py, create_python, function_name, examples)
    create_if_missing(problem_jl, create_julia, function_name, examples)
    create_if_missing(problem_rs, create_rust, function_name, examples)
    create_if_missing(problem_go, create_go, function_name, examples)