main.py

import csv
import argparse
from pathlib import Path
import sys


def read_csv(file_path: str, encoding: str = "utf-8") -> tuple[list[str], list[dict]]:
    path = Path(file_path)
    if not path.exists():
        print(f"Error: file not found: {file_path}", file=sys.stderr)
        sys.exit(1)

    try:
        with path.open("r", encoding=encoding, newline="") as f:
            reader = csv.DictReader(f)
            if reader.fieldnames is None:
                print("Error: CSV header not found.", file=sys.stderr)
                sys.exit(1)

            rows = list(reader)
            return reader.fieldnames, rows

    except UnicodeDecodeError:
        print(f"Error: failed to read as {encoding}.", file=sys.stderr)
        sys.exit(1)
    except Exception as e:
        print(f"Error: failed to read CSV: {e}", file=sys.stderr)
        sys.exit(1)


def read_template(file_path: str, encoding: str = "utf-8") -> str:
    path = Path(file_path)
    if not path.exists():
        print(f"Error: template file not found: {file_path}", file=sys.stderr)
        sys.exit(1)

    try:
        return path.read_text(encoding=encoding)
    except Exception as e:
        print(f"Error: failed to read template: {e}", file=sys.stderr)
        sys.exit(1)


def remove_empty_rows(rows: list[dict]) -> list[dict]:
    cleaned = []
    for row in rows:
        if any(str(v).strip() for v in row.values() if v is not None):
            cleaned.append(row)
    return cleaned


def filter_rows(rows: list[dict], keyword: str | None) -> list[dict]:
    if not keyword:
        return rows

    keyword_lower = keyword.lower()
    filtered = []

    for row in rows:
        row_text = " ".join(str(v) for v in row.values() if v is not None).lower()
        if keyword_lower in row_text:
            filtered.append(row)

    return filtered


def validate_columns(requested: list[str], available: list[str]) -> None:
    missing = [col for col in requested if col not in available]
    if missing:
        print(f"Error: column(s) not found: {', '.join(missing)}", file=sys.stderr)
        print(f"Available columns: {', '.join(available)}", file=sys.stderr)
        sys.exit(1)


def select_columns(rows: list[dict], columns: list[str]) -> list[dict]:
    return [{col: row.get(col, "") for col in columns} for row in rows]


def write_csv(file_path: str, columns: list[str], rows: list[dict], encoding: str = "utf-8") -> None:
    try:
        with open(file_path, "w", encoding=encoding, newline="") as f:
            writer = csv.DictWriter(f, fieldnames=columns)
            writer.writeheader()
            writer.writerows(rows)
    except Exception as e:
        print(f"Error: failed to write CSV: {e}", file=sys.stderr)
        sys.exit(1)


def write_txt(file_path: str, columns: list[str], rows: list[dict], encoding: str = "utf-8") -> None:
    try:
        with open(file_path, "w", encoding=encoding) as f:
            f.write(",".join(columns) + "\n")
            for row in rows:
                line = ",".join(str(row.get(col, "")) for col in columns)
                f.write(line + "\n")
    except Exception as e:
        print(f"Error: failed to write text file: {e}", file=sys.stderr)
        sys.exit(1)


def print_rows(columns: list[str], rows: list[dict]) -> None:
    print(",".join(columns))
    for row in rows:
        print(",".join(str(row.get(col, "")) for col in columns))


def render_template(template_text: str, rows: list[dict]) -> list[str]:
    rendered_list = []

    for row in rows:
        try:
            rendered = template_text.format(**row)
            rendered_list.append(rendered)
        except KeyError as e:
            print(f"Error: template variable not found in CSV columns: {e}", file=sys.stderr)
            sys.exit(1)
        except Exception as e:
            print(f"Error: failed to render template: {e}", file=sys.stderr)
            sys.exit(1)

    return rendered_list


def output_rendered_text(rendered_list: list[str], output_path: str | None, encoding: str = "utf-8") -> None:
    content = "\n---\n".join(rendered_list)

    if output_path:
        try:
            with open(output_path, "w", encoding=encoding) as f:
                f.write(content)
        except Exception as e:
            print(f"Error: failed to write rendered text: {e}", file=sys.stderr)
            sys.exit(1)
    else:
        print(content)


def parse_args():
    parser = argparse.ArgumentParser(
        description="Extract columns and rows from a CSV file, or render them with a template."
    )
    parser.add_argument("input", help="Input CSV file path")
    parser.add_argument(
        "--columns",
        help="Comma-separated column names to extract (example: ip,port,vlan)",
    )
    parser.add_argument(
        "--keyword",
        help="Filter rows containing this keyword",
    )
    parser.add_argument(
        "--output",
        help="Output file path (.csv or .txt). If omitted, print to stdout.",
    )
    parser.add_argument(
        "--encoding",
        default="utf-8",
        help="File encoding (default: utf-8). Example: cp932",
    )
    parser.add_argument(
        "--template",
        help="Template text file path for rendering output",
    )
    return parser.parse_args()


def main():
    args = parse_args()

    headers, rows = read_csv(args.input, encoding=args.encoding)
    rows = remove_empty_rows(rows)
    rows = filter_rows(rows, args.keyword)

    if args.columns:
        columns = [col.strip() for col in args.columns.split(",") if col.strip()]
        validate_columns(columns, headers)
    else:
        columns = headers

    rows = select_columns(rows, columns)

    if args.template:
        template_text = read_template(args.template, encoding=args.encoding)
        rendered_list = render_template(template_text, rows)

        if args.output and not args.output.lower().endswith(".txt"):
            print("Error: template output must be a .txt file", file=sys.stderr)
            sys.exit(1)

        output_rendered_text(rendered_list, args.output, encoding=args.encoding)
        return

    if args.output:
        output_path = args.output.lower()
        if output_path.endswith(".csv"):
            write_csv(args.output, columns, rows, encoding=args.encoding)
        elif output_path.endswith(".txt"):
            write_txt(args.output, columns, rows, encoding=args.encoding)
        else:
            print("Error: output file must end with .csv or .txt", file=sys.stderr)
            sys.exit(1)
    else:
        print_rows(columns, rows)


if __name__ == "__main__":
    main()