#!/usr/bin/env python
"""Batch ingest all files of a specific modality from a folder.

Wraps the processing functions from ingest_multimodal.py and adds
modality/filetype filtering, progress reporting, and dependency pre-checks.

Usage:
    # All text files in a folder
    python ingest_folder.py --input /path/to/docs --modality text --output-dir output/

    # Only PDFs
    python ingest_folder.py --input /path/to/docs --filetype pdf --output-dir output/

    # All images with PaddleOCR
    python ingest_folder.py --input /path/to/images --modality image --ocr-engine paddle

    # Everything in a folder (no filter)
    python ingest_folder.py --input /path/to/data --output-dir output/
"""

from __future__ import annotations

import argparse
import json
import shutil
import sys
from collections import Counter
from datetime import datetime, timezone
from pathlib import Path
from typing import Dict, Iterable, List, Optional

# Ensure sibling module is importable
sys.path.insert(0, str(Path(__file__).resolve().parent))

from ingest_multimodal import (
    TEXT_EXTS,
    IMAGE_EXTS,
    AUDIO_EXTS,
    VIDEO_EXTS,
    detect_file_type,
    build_doc_id,
    add_relations,
    load_langchain_documents,
    records_from_documents,
    records_from_image,
    records_from_audio,
    records_from_csv,
    process_video_chunked,
    OcrRunner,
    WhisperTranscriber,
    _resolve_modality,
)

# ---------------------------------------------------------------------------
# Modality mapping
# ---------------------------------------------------------------------------

MODALITY_EXTS = {
    "text": TEXT_EXTS,
    "image": IMAGE_EXTS,
    "audio": AUDIO_EXTS,
    "video": VIDEO_EXTS,
    "table": {".csv"},
}


# ---------------------------------------------------------------------------
# Filtered file iterator
# ---------------------------------------------------------------------------

def iter_files_filtered(
    root: Path,
    modality: Optional[str] = None,
    filetype: Optional[str] = None,
    max_files: Optional[int] = None,
) -> Iterable[Path]:
    """Recursively discover files with optional modality/filetype filtering.

    Args:
        root: Root directory to search.
        modality: One of 'text', 'image', 'audio', 'video', 'table', or None for all.
        filetype: Specific extension (without dot) like 'pdf', 'png', 'flac'.
                  Overrides modality if both are provided.
        max_files: Optional cap on number of files yielded.
    """
    allowed_exts = None
    if filetype:
        ext = f".{filetype.lower().lstrip('.')}"
        allowed_exts = {ext}
    elif modality and modality != "all":
        allowed_exts = MODALITY_EXTS.get(modality)
        if allowed_exts is None:
            raise ValueError(
                f"Unknown modality '{modality}'. Choose from: {', '.join(MODALITY_EXTS)}"
            )

    count = 0
    for path in sorted(root.rglob("*")):
        if path.is_dir():
            continue
        if path.name.startswith("."):
            continue
        if "output" in path.parts:
            continue

        if allowed_exts is not None:
            if path.suffix.lower() not in allowed_exts:
                continue

        if detect_file_type(path) is None:
            continue

        if max_files is not None and count >= max_files:
            break
        count += 1
        yield path


# ---------------------------------------------------------------------------
# Dependency pre-checks
# ---------------------------------------------------------------------------

def check_dependencies(modality: Optional[str], filetype: Optional[str]) -> None:
    """Verify required system tools and Python packages are available."""
    needs_ocr = False
    needs_whisper = False
    needs_ffmpeg = False

    if filetype:
        ext = f".{filetype.lower().lstrip('.')}"
        if ext in IMAGE_EXTS:
            needs_ocr = True
        elif ext in AUDIO_EXTS:
            needs_whisper = True
        elif ext in VIDEO_EXTS:
            needs_whisper = True
            needs_ffmpeg = True
    elif modality == "image":
        needs_ocr = True
    elif modality == "audio":
        needs_whisper = True
    elif modality == "video":
        needs_whisper = True
        needs_ffmpeg = True
    elif modality in (None, "all"):
        # Can't pre-check for all, dependencies checked per-file at runtime
        return

    errors = []
    if needs_ocr:
        if shutil.which("tesseract") is None:
            try:
                import pytesseract  # noqa: F401
            except ImportError:
                errors.append("tesseract is not installed. Install with: brew install tesseract")

    if needs_whisper:
        try:
            import faster_whisper  # noqa: F401
        except ImportError:
            errors.append(
                "faster-whisper is not installed. "
                "Install with: pip install faster-whisper (requires Python <3.13)"
            )

    if needs_ffmpeg:
        if shutil.which("ffmpeg") is None:
            errors.append("ffmpeg is not installed. Install with: brew install ffmpeg")
        if shutil.which("ffprobe") is None:
            errors.append("ffprobe is not installed (usually bundled with ffmpeg).")

    if errors:
        print("Missing dependencies:")
        for err in errors:
            print(f"  - {err}")
        sys.exit(1)


# ---------------------------------------------------------------------------
# Main ingestion
# ---------------------------------------------------------------------------

def ingest_folder(
    input_dir: Path,
    output_dir: Path,
    output_name: str = "cleaned_documents.jsonl",
    errors_name: str = "ingestion_errors.jsonl",
    modality: Optional[str] = None,
    filetype: Optional[str] = None,
    pdf_strategy: str = "auto",
    ocr_engine: str = "tesseract",
    ocr_language: str = "eng",
    whisper_model: str = "small",
    whisper_device: str = "cpu",
    whisper_compute_type: str = "int8",
    max_files: Optional[int] = None,
    chunk_duration: float = 600.0,
    extract_keyframes: bool = False,
    scene_threshold: float = 0.3,
) -> Dict:
    """Ingest all matching files from a folder into downstream-ready JSONL."""
    output_dir.mkdir(parents=True, exist_ok=True)
    output_path = output_dir / output_name
    errors_path = output_dir / errors_name

    file_counts: Counter = Counter()
    record_counts: Counter = Counter()
    error_count = 0
    total_files = 0
    total_records = 0

    success_files: List[str] = []
    failed_files: List[Dict[str, str]] = []
    empty_files: List[Dict[str, str]] = []

    ocr_runner = OcrRunner(ocr_engine, ocr_language)
    transcriber = None

    def _rel(path: Path) -> str:
        try:
            return path.relative_to(input_dir).as_posix()
        except ValueError:
            return path.as_posix()

    def log_error(path: Path, stage: str, err: Exception,
                  mod: str = "", ft: str = "") -> None:
        nonlocal error_count
        error_count += 1
        entry: Dict[str, object] = {
            "path": str(path), "stage": stage, "error": str(err),
        }
        if mod:
            entry["modality"] = mod
        if ft:
            entry["filetype"] = ft
        error_handle.write(json.dumps(entry, ensure_ascii=True) + "\n")

    with open(output_path, "w", encoding="utf-8") as out_handle, \
         open(errors_path, "w", encoding="utf-8") as error_handle:

        for path in iter_files_filtered(input_dir, modality, filetype, max_files):
            ft = detect_file_type(path)
            if not ft:
                continue
            total_files += 1
            file_counts[ft] += 1
            mod = _resolve_modality(ft)
            doc_id = build_doc_id(path, input_dir)

            # Progress reporting
            if total_files % 100 == 0:
                print(f"  [{total_files}] Processing: {path.name}")

            try:
                records: List[Dict] = []
                if ft == "csv":
                    records = records_from_csv(path, doc_id, input_dir)
                elif ft == "image":
                    records = records_from_image(path, doc_id, input_dir, ocr_runner)
                elif ft == "audio":
                    if transcriber is None:
                        transcriber = WhisperTranscriber(
                            whisper_model, whisper_device, whisper_compute_type
                        )
                    records = records_from_audio(
                        path, doc_id, input_dir, transcriber, "audio"
                    )
                elif ft == "video":
                    if transcriber is None:
                        transcriber = WhisperTranscriber(
                            whisper_model, whisper_device, whisper_compute_type
                        )
                    records = process_video_chunked(
                        video_path=path,
                        doc_id=doc_id,
                        root=input_dir,
                        transcriber=transcriber,
                        output_dir=output_dir,
                        chunk_duration=chunk_duration,
                        extract_keyframes=extract_keyframes,
                        ocr_runner=ocr_runner if extract_keyframes else None,
                        scene_threshold=scene_threshold,
                    )
                else:
                    docs, loader_name = load_langchain_documents(path, ft, pdf_strategy)
                    records = records_from_documents(
                        docs, doc_id, path, input_dir, ft, "text", loader_name
                    )

                if records:
                    # Video records already have relations from process_video_chunked
                    if ft != "video":
                        add_relations(records)
                    for rec in records:
                        out_handle.write(json.dumps(rec, ensure_ascii=True) + "\n")
                    total_records += len(records)
                    record_counts[ft] += len(records)
                    success_files.append(_rel(path))
                else:
                    empty_files.append({
                        "path": _rel(path),
                        "modality": mod,
                        "reason": "Extraction produced no content",
                    })
                    log_error(path, "empty_content",
                              Exception("Extraction produced no content"),
                              mod=mod, ft=ft)
            except Exception as exc:
                log_error(path, "ingest", exc, mod=mod, ft=ft)
                failed_files.append({
                    "path": _rel(path),
                    "modality": mod,
                    "reason": str(exc),
                })

    stats = {
        "input_root": str(input_dir),
        "output_file": str(output_path),
        "errors_file": str(errors_path),
        "generated_at": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
        "config": {
            "modality_filter": modality or "all",
            "filetype_filter": filetype,
            "pdf_strategy": pdf_strategy,
            "ocr_engine": ocr_engine,
            "ocr_language": ocr_language,
            "whisper_model": whisper_model,
            "whisper_device": whisper_device,
            "whisper_compute_type": whisper_compute_type,
            "video_chunk_duration": chunk_duration,
            "extract_keyframes": extract_keyframes,
            "scene_threshold": scene_threshold if extract_keyframes else None,
        },
        "counts": {
            "files_processed": total_files,
            "files_succeeded": len(success_files),
            "files_failed": len(failed_files),
            "files_empty": len(empty_files),
            "records_emitted": total_records,
            "errors": error_count,
        },
        "files_by_type": dict(file_counts),
        "records_by_type": dict(record_counts),
        "success_files": success_files,
        "failed_files": failed_files,
        "empty_files": empty_files,
    }

    stats_path = output_dir / "ingestion_stats.json"
    with open(stats_path, "w", encoding="utf-8") as f:
        json.dump(stats, f, indent=2, ensure_ascii=True)

    return stats


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Batch ingest all files of a specific type from a folder.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python ingest_folder.py --input /path/to/docs --modality text
  python ingest_folder.py --input /path/to/docs --filetype pdf
  python ingest_folder.py --input /path/to/images --modality image --ocr-engine paddle
  python ingest_folder.py --input /path/to/data --output-dir output/
        """,
    )
    parser.add_argument("--input", required=True, help="Root directory containing files to ingest.")
    parser.add_argument(
        "--output-dir",
        default="01_data_ingestion/output",
        help="Output directory for JSONL and stats (default: 01_data_ingestion/output).",
    )
    parser.add_argument("--output-name", default="cleaned_documents.jsonl", help="Output JSONL filename.")
    parser.add_argument(
        "--modality",
        choices=["text", "image", "audio", "video", "table", "all"],
        default="all",
        help="Filter by modality (default: all).",
    )
    parser.add_argument(
        "--filetype",
        default=None,
        help="Filter by specific extension (e.g. pdf, png, flac). Overrides --modality.",
    )
    parser.add_argument("--max-files", type=int, default=None, help="Optional cap on files processed.")
    parser.add_argument("--pdf-strategy", default="auto", help="Unstructured PDF strategy.")
    parser.add_argument("--ocr-engine", default="tesseract", choices=["tesseract", "paddle"])
    parser.add_argument("--ocr-language", default="eng", help="OCR language code.")
    parser.add_argument("--whisper-model", default="small", help="faster-whisper model size.")
    parser.add_argument("--whisper-device", default="cpu", help="faster-whisper device.")
    parser.add_argument("--whisper-compute-type", default="int8", help="faster-whisper compute type.")
    parser.add_argument("--chunk-duration", type=float, default=600.0,
                        help="Duration in seconds for video audio chunks (default: 600 = 10 min).")
    parser.add_argument("--extract-keyframes", action="store_true", default=False,
                        help="Extract keyframes from videos and run OCR on them.")
    parser.add_argument("--scene-threshold", type=float, default=0.3,
                        help="FFmpeg scene change threshold 0.0-1.0 (default: 0.3).")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    input_dir = Path(args.input)
    if not input_dir.exists():
        print(f"Error: input directory '{input_dir}' not found.")
        sys.exit(1)

    effective_modality = None if args.filetype else (args.modality if args.modality != "all" else None)
    check_dependencies(effective_modality, args.filetype)

    print(f"Ingesting from: {input_dir}")
    if args.filetype:
        print(f"  Filter: filetype={args.filetype}")
    elif args.modality != "all":
        print(f"  Filter: modality={args.modality}")
    else:
        print("  Filter: none (all supported types)")
    print(f"  Output: {args.output_dir}\n")

    stats = ingest_folder(
        input_dir=input_dir,
        output_dir=Path(args.output_dir),
        output_name=args.output_name,
        modality=effective_modality,
        filetype=args.filetype,
        pdf_strategy=args.pdf_strategy,
        ocr_engine=args.ocr_engine,
        ocr_language=args.ocr_language,
        whisper_model=args.whisper_model,
        whisper_device=args.whisper_device,
        whisper_compute_type=args.whisper_compute_type,
        max_files=args.max_files,
        chunk_duration=args.chunk_duration,
        extract_keyframes=args.extract_keyframes,
        scene_threshold=args.scene_threshold,
    )

    counts = stats["counts"]
    print(f"\nIngestion complete!")
    print(f"  Files processed:  {counts['files_processed']}")
    print(f"  Succeeded:        {counts['files_succeeded']}")
    print(f"  Failed:           {counts['files_failed']}")
    print(f"  Empty (no text):  {counts['files_empty']}")
    print(f"  Records emitted:  {counts['records_emitted']}")
    print(f"  Files by type:    {stats['files_by_type']}")
    print(f"  Records by type:  {stats['records_by_type']}")
    print(f"  Output: {stats['output_file']}")

    if counts["errors"] > 0:
        print(f"  Errors log: {stats['errors_file']}")


if __name__ == "__main__":
    main()