from __future__ import annotations

import argparse
import json
import re
from pathlib import Path
from typing import Dict, Iterable, List, Set, Tuple

DEFAULT_PROJECT_IDS: Tuple[str, str] = (
    "0f9c3e22-f4bb-4803-825a-4f238eeb06d5",
    "f3e02e84-e81c-4aa6-88d4-f5aa108c8227",
)

CHILD_NODE_TYPES: Tuple[str, ...] = (
    "contract",
    "contract_change",
    "construction_task",
    "construction_task_detail",
    "material",
    "material_detail",
    "contract_settlement",
    "contract_settlement_detail",
    "personnel",
    "personnel_detail",
    "preparation",
    "production",
    "lar",
    "lar_change",
    "project_detail_data",
    "units",
)

REQUIRED_NODE_TYPES: Tuple[str, ...] = ("project", "company", "supplier", *CHILD_NODE_TYPES)

FILENAME_PATTERN = re.compile(r"^(?P<node_type>.+)_\d+\.jsonl$")

SOURCE_NODE_ALIASES: Dict[str, Tuple[str, ...]] = {
    "units": ("units", "generator_unit"),
    "project_detail_data": ("project_detail_data", "project_detail"),
}


def _index_source_files(source_dir: Path) -> Dict[str, List[Path]]:
    index: Dict[str, List[Path]] = {}
    for file_path in sorted(source_dir.glob("*.jsonl")):
        match = FILENAME_PATTERN.match(file_path.name)
        if not match:
            continue
        node_type = match.group("node_type")
        index.setdefault(node_type, []).append(file_path)
    return index


def _read_jsonl(file_path: Path) -> List[dict]:
    rows: List[dict] = []
    with file_path.open("r", encoding="utf-8") as f:
        for line_no, line in enumerate(f, start=1):
            payload = line.strip()
            if not payload:
                continue
            try:
                rows.append(json.loads(payload))
            except json.JSONDecodeError as exc:
                raise ValueError(f"Invalid JSON in {file_path} at line {line_no}") from exc
    return rows


def _write_jsonl(file_path: Path, rows: Iterable[dict]) -> int:
    count = 0
    file_path.parent.mkdir(parents=True, exist_ok=True)
    with file_path.open("w", encoding="utf-8", newline="\n") as f:
        for row in rows:
            f.write(json.dumps(row, ensure_ascii=False))
            f.write("\n")
            count += 1
    return count


def _require_node_types(node_file_index: Dict[str, List[Path]], required: Iterable[str]) -> None:
    missing: List[str] = []
    for node in required:
        aliases = SOURCE_NODE_ALIASES.get(node, (node,))
        if not any(alias in node_file_index for alias in aliases):
            missing.append(node)
    if missing:
        raise FileNotFoundError(
            "Missing required datasource files for node types: "
            + ", ".join(sorted(missing))
            + ". Please ensure source datasource is complete."
        )


def _filter_rows_by_project(rows: List[dict], project_ids: Set[str]) -> List[dict]:
    return [row for row in rows if str(row.get("project_id", "")) in project_ids]


def _extract_contract_supplier_ids(contract_rows: List[dict]) -> Set[str]:
    supplier_ids: Set[str] = set()
    for row in contract_rows:
        for field in ("supplier_id", "company_id"):
            value = str(row.get(field, "")).strip()
            if value:
                supplier_ids.add(value)
                break
    return supplier_ids


def _supplier_identity_set(supplier_rows: List[dict]) -> Set[str]:
    identities: Set[str] = set()
    for row in supplier_rows:
        for field in ("_id", "id"):
            value = str(row.get(field, "")).strip()
            if value:
                identities.add(value)
    return identities


def _load_node_rows(node_file_index: Dict[str, List[Path]], node_type: str) -> List[dict]:
    aliases = SOURCE_NODE_ALIASES.get(node_type, (node_type,))
    source_node_type = next((alias for alias in aliases if alias in node_file_index), None)
    if source_node_type is None:
        raise FileNotFoundError(f"No source files found for node type: {node_type}, aliases={aliases}")

    rows: List[dict] = []
    for file_path in node_file_index[source_node_type]:
        rows.extend(_read_jsonl(file_path))
    return rows


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Filter filtered_datasource by project scope and write compact datasource for local testing."
    )
    parser.add_argument(
        "--source-dir",
        default="working_local_datasource/20260227-020003",
        help="Source datasource directory that contains *.jsonl files.",
    )
    parser.add_argument(
        "--output-dir",
        default="working_local_datasource/filtered",
        help="Output datasource directory (same JSONL format).",
    )
    parser.add_argument(
        "--project-ids",
        nargs="+",
        default=list(DEFAULT_PROJECT_IDS),
        help="Target project IDs to keep.",
    )
    parser.add_argument(
        "--clean-output",
        action="store_true",
        help="Delete existing *.jsonl under output-dir before writing.",
    )
    args = parser.parse_args()

    source_dir = Path(args.source_dir).resolve()
    output_dir = Path(args.output_dir).resolve()
    project_ids = {str(project_id).strip() for project_id in args.project_ids if str(project_id).strip()}

    if not source_dir.exists() or not source_dir.is_dir():
        raise FileNotFoundError(f"Source directory not found: {source_dir}")
    if not project_ids:
        raise ValueError("project-ids cannot be empty")

    node_file_index = _index_source_files(source_dir)
    _require_node_types(node_file_index, REQUIRED_NODE_TYPES)

    if args.clean_output and output_dir.exists():
        for stale in output_dir.glob("*.jsonl"):
            stale.unlink()

    project_rows = _load_node_rows(node_file_index, "project")
    selected_projects = [row for row in project_rows if str(row.get("_id", "")) in project_ids]
    selected_project_ids = {str(row.get("_id", "")) for row in selected_projects if str(row.get("_id", ""))}

    missing_projects = sorted(project_ids - selected_project_ids)
    if missing_projects:
        raise ValueError(f"Project IDs not found in source project data: {missing_projects}")

    filtered_by_node: Dict[str, List[dict]] = {"project": selected_projects}

    for node_type in CHILD_NODE_TYPES:
        rows = _load_node_rows(node_file_index, node_type)
        filtered_by_node[node_type] = _filter_rows_by_project(rows, selected_project_ids)

    selected_company_ids = {
        str(row.get("company_id", ""))
        for row in selected_projects
        if str(row.get("company_id", ""))
    }

    company_rows = _load_node_rows(node_file_index, "company")
    filtered_by_node["company"] = [
        row for row in company_rows if str(row.get("_id", "")) in selected_company_ids
    ]

    selected_supplier_ids = _extract_contract_supplier_ids(filtered_by_node["contract"])

    supplier_rows = _load_node_rows(node_file_index, "supplier")
    supplier_identity_set = _supplier_identity_set(supplier_rows)

    missing_supplier_ids = sorted(selected_supplier_ids - supplier_identity_set)
    if missing_supplier_ids:
        sample = missing_supplier_ids[:10]
        raise ValueError(
            "Contract supplier references missing in supplier datasource: "
            f"missing={len(missing_supplier_ids)}, sample={sample}"
        )

    filtered_by_node["supplier"] = [
        row
        for row in supplier_rows
        if str(row.get("_id", "")).strip() in selected_supplier_ids
        or str(row.get("id", "")).strip() in selected_supplier_ids
    ]

    write_order = ["project", *CHILD_NODE_TYPES, "company", "supplier"]
    output_counts: Dict[str, int] = {}

    for node_type in write_order:
        out_file = output_dir / f"{node_type}_1.jsonl"
        output_counts[node_type] = _write_jsonl(out_file, filtered_by_node[node_type])

    print("[Done] output_dir:", output_dir)
    print("[Done] project_ids:", sorted(selected_project_ids))
    for node_type in write_order:
        print(f"[Count] {node_type}: {output_counts[node_type]}")


if __name__ == "__main__":
    main()
