from __future__ import annotations

import argparse
import csv
import json
import re
import sys
from pathlib import Path

try:
    from openpyxl import load_workbook
except ImportError:
    load_workbook = None


FILE = Path(__file__).resolve()
DEFAULT_INPUT_FILE = FILE.parents[1] / "examples" / "cncap" / "G1M3_AFS1616_CNCAP-2024_11月_0306.xlsx"
DEFAULT_COLUMN_NAME = "原始数据地址"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Extract one column from a CSV/XLSX table and print or save the values."
    )
    parser.add_argument(
        "--input-file",
        type=str,
        default=str(DEFAULT_INPUT_FILE),
        help="Path to the input .xlsx or .csv file.",
    )
    parser.add_argument(
        "--column-name",
        type=str,
        default=DEFAULT_COLUMN_NAME,
        help="Header name of the target column.",
    )
    parser.add_argument(
        "--sheet-name",
        type=str,
        default="",
        help="Worksheet name for .xlsx files. Defaults to the first sheet.",
    )
    parser.add_argument(
        "--output-file",
        type=str,
        default="",
        help="Optional output text file. If omitted, values are written to stdout.",
    )
    parser.add_argument(
        "--json-file",
        type=str,
        default="",
        help="Optional output json file. Defaults to a sibling json file next to the input table.",
    )
    parser.add_argument(
        "--dedupe",
        action="store_true",
        help="Remove duplicate values while preserving the original order.",
    )
    parser.add_argument(
        "--list-columns",
        action="store_true",
        help="List the discovered header names and exit.",
    )
    return parser.parse_args()


def normalize_header(value: str) -> str:
    return re.sub(r"\s+", "", str(value or "")).strip()


def sanitize_filename(value: str) -> str:
    sanitized = re.sub(r'[\\/:*?"<>|]+', "_", str(value or "").strip())
    sanitized = re.sub(r"\s+", "_", sanitized)
    return sanitized.strip("._") or "column"


def build_default_json_path(input_path: Path, column_name: str) -> Path:
    return input_path.with_name(f"{input_path.stem}_{sanitize_filename(column_name)}.json")


def load_xlsx_table(path: Path, sheet_name: str) -> tuple[list[str], list[dict[str, str]], str]:
    if load_workbook is None:
        raise ImportError("openpyxl is required for .xlsx files. Please install it first.")

    workbook = load_workbook(path, read_only=True, data_only=True)
    try:
        if sheet_name:
            if sheet_name not in workbook.sheetnames:
                available = ", ".join(workbook.sheetnames)
                raise ValueError(f"Worksheet {sheet_name!r} not found. Available sheets: {available}")
            worksheet = workbook[sheet_name]
        else:
            worksheet = workbook[workbook.sheetnames[0]]

        header_row_values: list[str] | None = None
        header_indices: list[int] = []
        records: list[dict[str, str]] = []

        for row in worksheet.iter_rows(values_only=True):
            normalized_row = [str(value).strip() if value is not None else "" for value in row]
            if not any(normalized_row):
                continue

            if header_row_values is None:
                header_indices = [index for index, value in enumerate(normalized_row) if value]
                header_row_values = [normalized_row[index] for index in header_indices]
                continue

            record = {
                header_row_values[index]: normalized_row[col_idx] if col_idx < len(normalized_row) else ""
                for index, col_idx in enumerate(header_indices)
            }
            if any(record.values()):
                records.append(record)

        if header_row_values is None:
            raise ValueError("The worksheet is empty.")

        return header_row_values, records, worksheet.title
    finally:
        workbook.close()


def load_csv_table(path: Path) -> tuple[list[str], list[dict[str, str]], str]:
    with path.open("r", encoding="utf-8-sig", newline="") as file:
        sample = file.read(4096)
        file.seek(0)
        dialect = csv.Sniffer().sniff(sample) if sample.strip() else csv.excel
        reader = csv.DictReader(file, dialect=dialect)
        headers = reader.fieldnames or []
        records = []
        for row in reader:
            normalized_row = {str(key).strip(): str(value or "").strip() for key, value in row.items() if key is not None}
            if any(normalized_row.values()):
                records.append(normalized_row)
    return [str(header).strip() for header in headers], records, ""


def load_table(path: Path, sheet_name: str) -> tuple[list[str], list[dict[str, str]], str]:
    suffix = path.suffix.lower()
    if suffix == ".xlsx":
        return load_xlsx_table(path, sheet_name)
    if suffix == ".csv":
        return load_csv_table(path)
    raise ValueError(f"Unsupported input format: {suffix}. Only .xlsx and .csv are supported.")


def extract_column(records: list[dict[str, str]], column_name: str, dedupe: bool) -> list[str]:
    if not records:
        return []

    normalized_to_actual = {normalize_header(name): name for name in records[0].keys()}
    target_key = normalized_to_actual.get(normalize_header(column_name))
    if target_key is None:
        available = ", ".join(records[0].keys())
        raise ValueError(f"Column {column_name!r} not found. Available columns: {available}")

    values = [record.get(target_key, "").strip() for record in records]
    values = [value for value in values if value]
    if not dedupe:
        return values

    deduped_values: list[str] = []
    seen: set[str] = set()
    for value in values:
        if value in seen:
            continue
        seen.add(value)
        deduped_values.append(value)
    return deduped_values


def write_output(values: list[str], output_file: str) -> None:
    content = "\n".join(values)
    if output_file:
        output_path = Path(output_file)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        output_path.write_text(content + ("\n" if values else ""), encoding="utf-8")
        print(f"Saved {len(values)} rows to {output_path}", file=sys.stderr)
        return

    if content:
        sys.stdout.write(content)
        sys.stdout.write("\n")


def write_json_output(
    input_path: Path,
    resolved_sheet_name: str,
    column_name: str,
    values: list[str],
    json_file: str,
) -> Path:
    json_path = Path(json_file) if json_file else build_default_json_path(input_path, column_name)
    json_path.parent.mkdir(parents=True, exist_ok=True)
    payload = {
        "input_file": str(input_path),
        "sheet_name": resolved_sheet_name,
        "column_name": column_name,
        "num_rows": len(values),
        "values": values,
    }
    json_path.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    print(f"Saved json to {json_path}", file=sys.stderr)
    return json_path


def main() -> int:
    args = parse_args()
    input_path = Path(args.input_file)
    if not input_path.is_file():
        raise FileNotFoundError(f"Input file not found: {input_path}")

    headers, records, resolved_sheet_name = load_table(input_path, args.sheet_name)
    if args.list_columns:
        for header in headers:
            print(header)
        return 0

    values = extract_column(records, args.column_name, args.dedupe)
    write_output(values, args.output_file)
    json_path = write_json_output(input_path, resolved_sheet_name, args.column_name, values, args.json_file)

    sheet_info = f", sheet={resolved_sheet_name}" if resolved_sheet_name else ""
    print(
        f"Extracted {len(values)} rows from column {args.column_name!r} in {input_path}{sheet_info}, json={json_path}",
        file=sys.stderr,
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
