from __future__ import annotations

import argparse
import json
import re
from pathlib import Path
from typing import Any, Dict, List, Optional

from db.brochure_pdf_taxonomy import enrich_payload_with_pdf_taxonomy
from db.collection_brochure_seed import load_collection_brochure_seed

PDF_FALLBACK_ALIASES = {
    "johnsonwhite_site": "JohnsonSandstone_0226.pdf",
    "merrillslate_site": "MerrillPolarWhite_0226.pdf",
    "questmetrofrost_site": "FabuwoodMetroFrost_0326.pdf",
}


def enrich_brochure_file(
    *,
    json_path: str,
    pdf_folder: str,
    output_folder: str,
) -> Dict[str, Any]:
    payload = load_collection_brochure_seed(json_path)
    pdf_path = _match_pdf_path(payload, pdf_folder=pdf_folder, json_path=json_path)
    enriched = enrich_payload_with_pdf_taxonomy(payload, pdf_path=pdf_path)

    out_dir = Path(output_folder)
    out_dir.mkdir(parents=True, exist_ok=True)
    out_path = out_dir / Path(json_path).name
    out_path.write_text(json.dumps(enriched, indent=2), encoding="utf-8")
    return {
        "json_path": json_path,
        "pdf_path": str(pdf_path),
        "output_path": str(out_path),
        "group_count": len(enriched.get("taxonomy_groups") or []),
    }


def enrich_brochure_folder(
    *,
    json_folder: str,
    pdf_folder: str,
    output_folder: str,
    skip_tokens: Optional[List[str]] = None,
) -> Dict[str, Any]:
    folder = Path(json_folder)
    paths = sorted(path for path in folder.glob("*.json") if path.is_file())
    skip_set = {token.strip().lower() for token in (skip_tokens or []) if token and token.strip()}
    selected = [path for path in paths if path.name.lower() not in skip_set and path.stem.lower() not in skip_set]
    results: List[Dict[str, Any]] = []
    errors: List[Dict[str, Any]] = []
    for path in selected:
        try:
            results.append(
                enrich_brochure_file(
                    json_path=str(path),
                    pdf_folder=pdf_folder,
                    output_folder=output_folder,
                )
            )
        except Exception as exc:
            errors.append({"json_path": str(path), "error": str(exc)})
    return {
        "file_count": len(selected),
        "enriched_count": len(results),
        "error_count": len(errors),
        "results": results,
        "errors": errors,
    }


def _match_pdf_path(payload: Dict[str, Any], *, pdf_folder: str, json_path: str) -> Path:
    pdf_dir = Path(pdf_folder)
    available_pdfs = [path for path in pdf_dir.glob("*.pdf") if path.is_file()]
    source_document = payload.get("source_document") if isinstance(payload.get("source_document"), dict) else {}
    candidate_names: List[str] = []
    file_name = str(source_document.get("file_name") or "").strip()
    if file_name:
        candidate_names.append(Path(file_name).with_suffix(".pdf").name)
    json_file_name = Path(json_path).with_suffix(".pdf").name
    json_stem = Path(json_path).stem
    candidate_names.append(json_file_name)

    fallback_name = PDF_FALLBACK_ALIASES.get(_normalize_name(json_stem))
    if fallback_name:
        candidate_names.append(fallback_name)

    collection = payload.get("collection") if isinstance(payload.get("collection"), dict) else {}
    for value in (
        collection.get("display_name"),
        collection.get("line_name"),
        collection.get("canonical_code"),
        collection.get("color_label"),
        source_document.get("collection_code"),
    ):
        text = str(value or "").strip()
        if not text:
            continue
        flat = "".join(ch for ch in text if ch.isalnum())
        if flat:
            candidate_names.extend(path.name for path in available_pdfs if _normalize_name(flat) in _normalize_name(path.stem))

    seen = set()
    for name in candidate_names:
        if not name or name in seen:
            continue
        seen.add(name)
        path = pdf_dir / name
        if path.is_file():
            return path

    fuzzy_match = _best_fuzzy_pdf_match(
        payload=payload,
        json_path=json_path,
        available_pdfs=available_pdfs,
    )
    if fuzzy_match is not None:
        return fuzzy_match
    raise FileNotFoundError(f"No matching PDF found for {json_path}")


def _best_fuzzy_pdf_match(*, payload: Dict[str, Any], json_path: str, available_pdfs: List[Path]) -> Optional[Path]:
    candidate_tokens: set[str] = set()
    source_document = payload.get("source_document") if isinstance(payload.get("source_document"), dict) else {}
    collection = payload.get("collection") if isinstance(payload.get("collection"), dict) else {}

    for value in (
        Path(json_path).stem,
        source_document.get("collection_code"),
        collection.get("display_name"),
        collection.get("line_name"),
        collection.get("canonical_code"),
        collection.get("color_label"),
    ):
        candidate_tokens.update(_tokenize_name(value))
    candidate_tokens.discard("site")
    candidate_tokens.discard("json")
    if not candidate_tokens:
        return None

    scored: List[tuple[int, int, Path]] = []
    for pdf_path in available_pdfs:
        pdf_tokens = _tokenize_name(pdf_path.stem)
        overlap = len(candidate_tokens & pdf_tokens)
        if overlap <= 0:
            continue
        scored.append((overlap, len(pdf_tokens), pdf_path))

    if not scored:
        return None
    scored.sort(key=lambda item: (item[0], -item[1], item[2].name), reverse=True)
    best_overlap = scored[0][0]
    top = [item for item in scored if item[0] == best_overlap]
    if len(top) == 1:
        return top[0][2]
    return None


def _normalize_name(value: Any) -> str:
    return "".join(ch.lower() for ch in str(value or "") if ch.isalnum())


def _tokenize_name(value: Any) -> set[str]:
    raw = str(value or "").strip()
    if not raw:
        return set()
    pieces = re.split(r"[^A-Za-z0-9]+", raw)
    expanded: List[str] = []
    for piece in pieces:
        if not piece:
            continue
        expanded.extend(part for part in re.findall(r"[A-Z]+(?=[A-Z][a-z]|\d|$)|[A-Z]?[a-z]+|\d+", piece) if part)
    return {part.lower() for part in expanded if part}


def main() -> None:
    parser = argparse.ArgumentParser(description="Enrich brochure JSON files with explicit taxonomy_groups from PDFs.")
    parser.add_argument("--file", default=None, help="One brochure JSON file.")
    parser.add_argument("--json-folder", default=None, help="Folder containing brochure JSON files.")
    parser.add_argument("--pdf-folder", required=True, help="Folder containing brochure PDFs.")
    parser.add_argument("--output-folder", required=True, help="Folder to write enriched JSON files.")
    parser.add_argument("--skip", action="append", default=None, help="File stem/name to skip; repeatable.")
    args = parser.parse_args()

    if args.file:
        result = enrich_brochure_file(
            json_path=args.file,
            pdf_folder=args.pdf_folder,
            output_folder=args.output_folder,
        )
    elif args.json_folder:
        result = enrich_brochure_folder(
            json_folder=args.json_folder,
            pdf_folder=args.pdf_folder,
            output_folder=args.output_folder,
            skip_tokens=args.skip,
        )
    else:
        raise SystemExit("Provide --file or --json-folder")
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()