from collections.abc import Awaitable, Callable, Mapping
from dataclasses import dataclass

from mirage.commands.builtin.utils.lines import split_lines
from mirage.commands.config import CommandOpts
from mirage.commands.spec import SPECS
from mirage.commands.spec.types import CommandName, FlagValue, FlagView
from mirage.commands.spec.usage import extra_operand_error
from mirage.io.types import ByteSource, IOResult
from mirage.types import PathSpec


def _split_fields(line: str, delimiter: str | None) -> list[str]:
    return line.split(delimiter) if delimiter else line.split()


def _field_at(fields: list[str], index: int) -> str:
    """Return a field by index, or the empty string when the line is short.

    Blank and short lines are legal input for GNU ``join``: the missing
    field simply reads as an empty key rather than aborting the run.

    Args:
        fields (list[str]): Fields of one input line.
        index (int): Zero-based field index to read.
    """
    return fields[index] if 0 <= index < len(fields) else ""


def _build_join_map(
    lines: list[str],
    field_idx: int,
    delimiter: str | None,
    ignore_case: bool,
) -> dict[str, list[list[str]]]:
    result: dict[str, list[list[str]]] = {}
    for line in lines:
        parts = _split_fields(line, delimiter)
        if field_idx < len(parts):
            key = parts[field_idx].casefold(
            ) if ignore_case else parts[field_idx]
            if key not in result:
                result[key] = []
            result[key].append(parts)
    return result


def _rest(fields: list[str], key_idx: int) -> list[str]:
    return fields[:key_idx] + fields[key_idx + 1:]


def _format_row(
    key: str,
    fields1: list[str],
    field1: int,
    fields2: list[str],
    field2: int,
    o_fmt: str | None,
    out_sep: str,
    empty_value: str | None,
) -> str:
    if o_fmt is None:
        return out_sep.join([key] + _rest(fields1, field1) +
                            _rest(fields2, field2))
    # -o FILENUM.FIELD indexes the original 1-based field (the join key
    # included), so map against the full field list, not the key-stripped
    # rest; a missing field uses the -e replacement (GNU), else empty.
    fields: list[str] = []
    for spec in o_fmt.split(","):
        spec = spec.strip()
        if spec == "0":
            fields.append(key)
            continue
        file_part, _, field_part = spec.partition(".")
        src = fields1 if file_part == "1" else fields2
        idx = int(field_part) - 1
        if 0 <= idx < len(src):
            fields.append(src[idx])
        else:
            fields.append(empty_value if empty_value is not None else "")
    return out_sep.join(fields)


def _join_lines(
    lines1: list[str],
    lines2: list[str],
    field1: int,
    field2: int,
    sep: str | None,
    also_unpairable: str | None,
    only_unpairable: str | None,
    empty_value: str | None,
    output_format: str | None,
    ignore_case: bool,
) -> list[str]:
    map2 = _build_join_map(lines2, field2, sep, ignore_case)
    out_sep = sep if sep else " "
    out_lines: list[str] = []
    matched_keys2: set[str] = set()

    for line in lines1:
        parts = _split_fields(line, sep)
        if field1 >= len(parts):
            continue
        key = parts[field1]
        lookup_key = key.casefold() if ignore_case else key
        if lookup_key in map2:
            matched_keys2.add(lookup_key)
            if only_unpairable is None:
                for fields2 in map2[lookup_key]:
                    out_lines.append(
                        _format_row(key, parts, field1, fields2, field2,
                                    output_format, out_sep, empty_value))
        elif only_unpairable == "1" or also_unpairable == "1":
            out_lines.append(
                _format_row(key, parts, field1, [], field2, output_format,
                            out_sep, empty_value))

    if also_unpairable == "2" or only_unpairable == "2":
        for line in lines2:
            parts = _split_fields(line, sep)
            if field2 >= len(parts):
                continue
            key = parts[field2]
            lookup_key = key.casefold() if ignore_case else key
            if lookup_key not in matched_keys2:
                out_lines.append(
                    _format_row(key, [], field1, parts, field2, output_format,
                                out_sep, empty_value))

    return out_lines


async def join_cmd(
    paths: list[PathSpec],
    *,
    read_bytes: Callable[..., Awaitable[bytes]],
    field1: int = 0,
    field2: int = 0,
    separator: str | None = None,
    also_unpairable: str | None = None,
    only_unpairable: str | None = None,
    empty_value: str | None = None,
    output_format: str | None = None,
    ignore_case: bool = False,
    zero_terminated: bool = False,
    check_order: bool = False,
    header: bool = False,
) -> tuple[ByteSource | None, IOResult]:
    if len(paths) > 2:
        raise extra_operand_error(CommandName.JOIN, paths[2].raw_path
                                  or paths[2].virtual)
    if len(paths) < 2:
        raise ValueError("join: requires two paths")
    data1 = (await read_bytes(paths[0])).decode(errors="replace")
    data2 = (await read_bytes(paths[1])).decode(errors="replace")
    lines1 = data1.rstrip("\0").split(
        "\0") if zero_terminated else split_lines(data1)
    lines2 = data2.rstrip("\0").split(
        "\0") if zero_terminated else split_lines(data2)
    header_lines: list[str] = []
    if header and lines1 and lines2:
        first1 = _split_fields(lines1.pop(0), separator)
        first2 = _split_fields(lines2.pop(0), separator)
        header_lines.append(
            _format_row(_field_at(first1, field1), first1, field1, first2,
                        field2, output_format, separator or " ", empty_value))
    key_fn = str.casefold if ignore_case else str
    stderr = ""
    if check_order:
        keys1 = [
            key_fn(_field_at(_split_fields(line, separator), field1))
            for line in lines1
        ]
        keys2 = [
            key_fn(_field_at(_split_fields(line, separator), field2))
            for line in lines2
        ]
        if keys1 != sorted(keys1):
            stderr = "join: file 1 is not in sorted order\n"
        elif keys2 != sorted(keys2):
            stderr = "join: file 2 is not in sorted order\n"
    out_lines = _join_lines(lines1, lines2, field1, field2, separator,
                            also_unpairable, only_unpairable, empty_value,
                            output_format, ignore_case)
    out_lines = header_lines + out_lines
    if not out_lines:
        return None, IOResult(stderr=stderr.encode() if stderr else None,
                              exit_code=1 if stderr else 0)
    record_separator = "\0" if zero_terminated else "\n"
    return (record_separator.join(out_lines) +
            record_separator).encode(), IOResult(
                stderr=stderr.encode() if stderr else None,
                exit_code=1 if stderr else 0,
            )


__all__ = ["join_cmd"]


@dataclass(frozen=True, slots=True)
class JoinFlags:
    field1: int = 0
    field2: int = 0
    separator: str | None = None
    also_unpairable: str | None = None
    only_unpairable: str | None = None
    empty_value: str | None = None
    output_format: str | None = None
    ignore_case: bool = False
    zero_terminated: bool = False
    check_order: bool = False
    header: bool = False


def parse_flags(flags: Mapping[str, FlagValue]) -> JoinFlags:
    fl = FlagView(flags, spec=SPECS["join"])
    j = fl.as_str("j")
    return JoinFlags(
        field1=int(j or fl.as_str("args_1") or "1") - 1,
        field2=int(j or fl.as_str("2") or "1") - 1,
        separator=fl.as_str("t"),
        also_unpairable=fl.as_str("a"),
        only_unpairable=fl.as_str("v"),
        empty_value=fl.as_str("e"),
        output_format=fl.as_str("o"),
        ignore_case=fl.as_bool("ignore_case"),
        zero_terminated=fl.as_bool("zero_terminated"),
        check_order=fl.as_bool("check_order")
        and not fl.as_bool("nocheck_order"),
        header=fl.as_bool("header"),
    )


async def join_generic(
    paths: list[PathSpec],
    texts: list[str],
    opts: CommandOpts,
    read_bytes: Callable[..., Awaitable[bytes]],
) -> tuple[ByteSource | None, IOResult]:
    parsed = parse_flags(opts.flags)
    return await join_cmd(paths,
                          read_bytes=read_bytes,
                          field1=parsed.field1,
                          field2=parsed.field2,
                          separator=parsed.separator,
                          also_unpairable=parsed.also_unpairable,
                          only_unpairable=parsed.only_unpairable,
                          empty_value=parsed.empty_value,
                          output_format=parsed.output_format,
                          ignore_case=parsed.ignore_case,
                          zero_terminated=parsed.zero_terminated,
                          check_order=parsed.check_order,
                          header=parsed.header)
