# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Email: Georgescottfoley@proton.me
# Licensed under the MIT License - see LICENSE file for details

"""
File Reader Skill

Extract text from PDF, DOCX, XLSX, CSV, TXT, and images (OCR).
Auto-detects format and returns clean text for LLM processing.

Dependencies (all optional, graceful degradation):
  - PyMuPDF/fitz  (PDF text + table extraction)
  - python-docx   (DOCX reading)
  - openpyxl      (XLSX reading)
  - pytesseract   (OCR, requires Tesseract system package)
  - Pillow        (Image preprocessing for OCR)

Stdlib fallbacks:
  - csv           (CSV reading)
  - json          (JSON reading)
"""

import csv
import json
import logging
from pathlib import Path
from typing import Optional

logger = logging.getLogger(__name__)

# --- Optional dependency probing (auto-install on first use) ---

FITZ_AVAILABLE = False
DOCX_AVAILABLE = False
OPENPYXL_AVAILABLE = False
OCR_AVAILABLE = False

try:
    import fitz  # PyMuPDF

    FITZ_AVAILABLE = True
except ImportError as e:
    logger.debug("Optional dependency not available: %s", e)
try:
    from docx import Document as DocxDocument

    DOCX_AVAILABLE = True
except ImportError as e:
    logger.debug("Optional dependency not available: %s", e)
try:
    import openpyxl

    OPENPYXL_AVAILABLE = True
except ImportError as e:
    logger.debug("Optional dependency not available: %s", e)
try:
    import pytesseract
    from PIL import Image

    OCR_AVAILABLE = True
except ImportError as e:
    logger.debug("Optional dependency not available: %s", e)
def _ensure_fitz():
    global FITZ_AVAILABLE, fitz
    if FITZ_AVAILABLE:
        return True
    from familiar.core.deps import PDF_PACKAGES, ensure_packages

    ok, _ = ensure_packages(PDF_PACKAGES)
    if ok:
        import fitz as _fitz

        globals()["fitz"] = _fitz
        FITZ_AVAILABLE = True
    return FITZ_AVAILABLE


def _ensure_docx():
    global DOCX_AVAILABLE, DocxDocument
    if DOCX_AVAILABLE:
        return True
    from familiar.core.deps import DOCX_PACKAGES, ensure_packages

    ok, _ = ensure_packages(DOCX_PACKAGES)
    if ok:
        from docx import Document as _DocxDocument

        globals()["DocxDocument"] = _DocxDocument
        DOCX_AVAILABLE = True
    return DOCX_AVAILABLE


def _ensure_openpyxl():
    global OPENPYXL_AVAILABLE, openpyxl
    if OPENPYXL_AVAILABLE:
        return True
    from familiar.core.deps import XLSX_PACKAGES, ensure_packages

    ok, _ = ensure_packages(XLSX_PACKAGES)
    if ok:
        import openpyxl as _openpyxl

        globals()["openpyxl"] = _openpyxl
        OPENPYXL_AVAILABLE = True
    return OPENPYXL_AVAILABLE


def _ensure_ocr():
    global OCR_AVAILABLE, pytesseract, Image
    if OCR_AVAILABLE:
        return True
    from familiar.core.deps import OCR_PACKAGES, ensure_packages

    ok, _ = ensure_packages(OCR_PACKAGES)
    if ok:
        import pytesseract as _pytesseract
        from PIL import Image as _Image

        globals()["pytesseract"] = _pytesseract
        globals()["Image"] = _Image
        OCR_AVAILABLE = True
    return OCR_AVAILABLE


MAX_TEXT_LENGTH = 50000  # 50K chars max output to avoid overloading context


def _detect_format(filepath: Path) -> str:
    """Detect file format from extension."""
    ext = filepath.suffix.lower()
    format_map = {
        ".pdf": "pdf",
        ".docx": "docx",
        ".doc": "doc",
        ".xlsx": "xlsx",
        ".xls": "xls",
        ".csv": "csv",
        ".tsv": "tsv",
        ".txt": "txt",
        ".md": "txt",
        ".json": "json",
        ".xml": "txt",
        ".html": "txt",
        ".htm": "txt",
        ".rtf": "txt",
        ".png": "image",
        ".jpg": "image",
        ".jpeg": "image",
        ".tiff": "image",
        ".tif": "image",
        ".bmp": "image",
        ".gif": "image",
        ".webp": "image",
    }
    return format_map.get(ext, "unknown")


def _read_pdf(filepath: Path, pages: Optional[list] = None) -> str:
    """Extract text from PDF with page-level control."""
    if not _ensure_fitz():
        return "⚠️ PDF reading unavailable — PyMuPDF could not be installed."

    doc = fitz.open(str(filepath))
    texts = []
    total_pages = len(doc)

    page_range = pages if pages else range(total_pages)
    for page_num in page_range:
        if 0 <= page_num < total_pages:
            page = doc[page_num]
            text = page.get_text("text")
            if text.strip():
                texts.append(f"--- Page {page_num + 1} ---\n{text.strip()}")

    doc.close()

    if not texts:
        return f"No text extracted from PDF ({total_pages} pages). The document may be scanned — try ocr_image."

    result = "\n\n".join(texts)
    return result[:MAX_TEXT_LENGTH]


def _read_docx(filepath: Path) -> str:
    """Extract text from DOCX."""
    if not _ensure_docx():
        return "⚠️ DOCX reading unavailable — python-docx could not be installed."

    doc = DocxDocument(str(filepath))
    paragraphs = []
    for para in doc.paragraphs:
        if para.text.strip():
            prefix = ""
            if para.style.name.startswith("Heading"):
                try:
                    level = int(para.style.name.replace("Heading ", "").replace("Heading", "1"))
                    prefix = "#" * level + " "
                except ValueError:
                    prefix = "## "
            paragraphs.append(f"{prefix}{para.text}")

    # Also extract tables
    for i, table in enumerate(doc.tables):
        paragraphs.append(f"\n[Table {i + 1}]")
        for row in table.rows:
            cells = [cell.text.strip() for cell in row.cells]
            paragraphs.append(" | ".join(cells))

    result = "\n\n".join(paragraphs) if paragraphs else "No text content found in DOCX."
    return result[:MAX_TEXT_LENGTH]


def _read_xlsx(
    filepath: Path, sheet: Optional[str] = None, cell_range: Optional[str] = None
) -> str:
    """Read XLSX with sheet selection and range extraction."""
    if not _ensure_openpyxl():
        return "⚠️ XLSX reading unavailable — openpyxl could not be installed."

    wb = openpyxl.load_workbook(str(filepath), read_only=True, data_only=True)
    sheets = wb.sheetnames

    if sheet:
        if sheet not in sheets:
            wb.close()
            return f"Sheet '{sheet}' not found. Available sheets: {', '.join(sheets)}"
        ws = wb[sheet]
    else:
        ws = wb.active

    rows = []
    if cell_range:
        # Parse range like "A1:D10"
        for row in ws[cell_range]:
            cells = [str(cell.value) if cell.value is not None else "" for cell in row]
            rows.append(cells)
    else:
        for row in ws.iter_rows(values_only=True):
            cells = [str(v) if v is not None else "" for v in row]
            if any(c.strip() for c in cells):
                rows.append(cells)

    wb.close()

    if not rows:
        return f"No data found in sheet '{ws.title}'."

    # Detect headers (first row)
    lines = [f"Sheet: {ws.title} ({len(rows)} rows, {len(rows[0])} cols)\n"]

    # Format as aligned table
    for i, row in enumerate(rows[:200]):  # Cap at 200 rows
        line = " | ".join(c[:40] for c in row)
        lines.append(line)
        if i == 0:
            lines.append("-" * min(len(line), 120))

    if len(rows) > 200:
        lines.append(f"\n... and {len(rows) - 200} more rows")

    return "\n".join(lines)


def _read_csv(filepath: Path, delimiter: str = ",") -> str:
    """Read CSV/TSV file."""
    rows = []
    try:
        with open(filepath, encoding="utf-8", errors="replace", newline="") as f:
            # Sniff delimiter if not explicit
            sample = f.read(4096)
            f.seek(0)
            if delimiter == "," and "\t" in sample and sample.count("\t") > sample.count(","):
                delimiter = "\t"

            reader = csv.reader(f, delimiter=delimiter)
            for row in reader:
                if any(c.strip() for c in row):
                    rows.append(row)
                if len(rows) > 500:
                    break
    except Exception as e:
        return f"Error reading CSV: {e}"

    if not rows:
        return "No data found in file."

    lines = [f"CSV: {len(rows)} rows, {len(rows[0])} cols\n"]
    for i, row in enumerate(rows[:200]):
        line = " | ".join(c[:40] for c in row)
        lines.append(line)
        if i == 0:
            lines.append("-" * min(len(line), 120))

    if len(rows) > 200:
        lines.append(f"\n... and {len(rows) - 200} more rows")

    return "\n".join(lines)


def _read_text(filepath: Path) -> str:
    """Read plain text/markdown/html files."""
    try:
        text = filepath.read_text(encoding="utf-8", errors="replace")
        return text[:MAX_TEXT_LENGTH]
    except Exception as e:
        return f"Error reading text file: {e}"


def _ocr_image(filepath: Path) -> str:
    """Extract text from image via OCR."""
    if not _ensure_ocr():
        return "⚠️ OCR unavailable — pytesseract/Pillow could not be installed (also requires Tesseract system package)."

    try:
        image = Image.open(str(filepath))
        # Preprocess: convert to grayscale for better OCR
        if image.mode != "L":
            image = image.convert("L")
        text = pytesseract.image_to_string(image)
        if not text.strip():
            return "OCR completed but no text was detected in the image."
        return text[:MAX_TEXT_LENGTH]
    except Exception as e:
        return f"OCR error: {e}"


def _extract_pdf_tables(filepath: Path) -> list:
    """Extract tabular data from PDF into structured list of tables."""
    if not FITZ_AVAILABLE:
        return []

    doc = fitz.open(str(filepath))
    tables = []

    for page_num in range(len(doc)):
        page = doc[page_num]
        # Use PyMuPDF's built-in table detection (available in recent versions)
        try:
            page_tables = page.find_tables()
            for tab in page_tables:
                rows = []
                for row in tab.extract():
                    rows.append([str(cell) if cell else "" for cell in row])
                if rows:
                    tables.append(
                        {
                            "page": page_num + 1,
                            "headers": rows[0] if rows else [],
                            "rows": rows[1:] if len(rows) > 1 else [],
                        }
                    )
        except AttributeError:
            # Older PyMuPDF without find_tables — fall back to text blocks
            blocks = page.get_text("blocks")
            # Simple heuristic: look for aligned text blocks
            if blocks:
                text_rows = []
                for b in blocks:
                    if b[6] == 0:  # text block type
                        text_rows.append(b[4].strip())
                if text_rows:
                    tables.append(
                        {
                            "page": page_num + 1,
                            "headers": [],
                            "rows": [[r] for r in text_rows],
                            "note": "Extracted as text blocks (table detection unavailable)",
                        }
                    )

    doc.close()
    return tables


# === Tool Handlers ===


def read_document(data: dict) -> str:
    """Extract text from PDF, DOCX, XLSX, CSV, TXT. Auto-detects format."""
    filepath = data.get("filepath", "").strip()
    if not filepath:
        return "Please provide a file path."

    path = Path(filepath).expanduser()
    if not path.exists():
        return f"File not found: {filepath}"

    fmt = data.get("format") or _detect_format(path)

    if fmt == "pdf":
        return _read_pdf(path)
    elif fmt == "docx":
        return _read_docx(path)
    elif fmt in ("xlsx", "xls"):
        return _read_xlsx(path)
    elif fmt in ("csv", "tsv"):
        delimiter = "\t" if fmt == "tsv" else ","
        return _read_csv(path, delimiter)
    elif fmt == "json":
        try:
            content = json.loads(path.read_text(encoding="utf-8"))
            return json.dumps(content, indent=2)[:MAX_TEXT_LENGTH]
        except Exception as e:
            return f"Error reading JSON: {e}"
    elif fmt == "txt":
        return _read_text(path)
    elif fmt == "image":
        return _ocr_image(path)
    elif fmt == "doc":
        return "⚠️ Legacy .doc format not supported. Please convert to .docx first."
    else:
        # Try reading as text
        try:
            return _read_text(path)
        except Exception:
            return f"Unsupported file format: {path.suffix}"


def read_pdf(data: dict) -> str:
    """Extract text and tables from PDF with page-level control."""
    filepath = data.get("filepath", "").strip()
    if not filepath:
        return "Please provide a file path."

    path = Path(filepath).expanduser()
    if not path.exists():
        return f"File not found: {filepath}"

    pages = data.get("pages")  # Optional list of 0-indexed page numbers
    if isinstance(pages, str):
        # Parse "1-5" or "1,3,5" style input (convert to 0-indexed)
        page_list = []
        for part in pages.split(","):
            part = part.strip()
            if "-" in part:
                start, end = part.split("-", 1)
                page_list.extend(range(int(start) - 1, int(end)))
            else:
                page_list.append(int(part) - 1)
        pages = page_list

    return _read_pdf(path, pages)


def read_spreadsheet(data: dict) -> str:
    """Read XLSX/CSV with sheet selection, range extraction, header detection."""
    filepath = data.get("filepath", "").strip()
    if not filepath:
        return "Please provide a file path."

    path = Path(filepath).expanduser()
    if not path.exists():
        return f"File not found: {filepath}"

    fmt = _detect_format(path)
    sheet = data.get("sheet")
    cell_range = data.get("range")

    if fmt in ("xlsx", "xls"):
        result = _read_xlsx(path, sheet=sheet, cell_range=cell_range)

        # If requested, also list all sheet names
        if data.get("list_sheets") and OPENPYXL_AVAILABLE:
            wb = openpyxl.load_workbook(str(path), read_only=True)
            result = f"Sheets: {', '.join(wb.sheetnames)}\n\n{result}"
            wb.close()

        return result

    elif fmt in ("csv", "tsv"):
        delimiter = "\t" if fmt == "tsv" else ","
        return _read_csv(path, delimiter)

    else:
        return f"Not a spreadsheet format: {path.suffix}. Supported: .xlsx, .xls, .csv, .tsv"


def summarize_file(data: dict) -> str:
    """Read a file and return extracted text for LLM summarization."""
    filepath = data.get("filepath", "").strip()
    if not filepath:
        return "Please provide a file path."

    path = Path(filepath).expanduser()
    if not path.exists():
        return f"File not found: {filepath}"

    # Extract text using read_document
    text = read_document({"filepath": filepath})

    if text.startswith("⚠️") or text.startswith("File not found") or text.startswith("Error"):
        return text

    # Truncate for summarization context
    max_chars = int(data.get("max_chars", 12000))
    if len(text) > max_chars:
        text = text[:max_chars] + "\n\n[Content truncated for summarization]"

    size_kb = path.stat().st_size / 1024
    return (
        f"📄 File: {path.name} ({size_kb:.1f} KB, {_detect_format(path)} format)\n"
        f"Please summarize the following content:\n\n{text}"
    )


def extract_tables(data: dict) -> str:
    """Extract tabular data from PDF or DOCX into structured JSON."""
    filepath = data.get("filepath", "").strip()
    if not filepath:
        return "Please provide a file path."

    path = Path(filepath).expanduser()
    if not path.exists():
        return f"File not found: {filepath}"

    fmt = _detect_format(path)

    tables = []

    if fmt == "pdf":
        tables = _extract_pdf_tables(path)
    elif fmt == "docx":
        if not _ensure_docx():
            return "⚠️ DOCX reading unavailable — python-docx could not be installed."
        doc = DocxDocument(str(path))
        for i, table in enumerate(doc.tables):
            rows = []
            for row in table.rows:
                cells = [cell.text.strip() for cell in row.cells]
                rows.append(cells)
            if rows:
                tables.append(
                    {
                        "table_index": i + 1,
                        "headers": rows[0],
                        "rows": rows[1:] if len(rows) > 1 else [],
                    }
                )
    elif fmt in ("xlsx", "xls"):
        # Return the spreadsheet content as a table
        return read_spreadsheet(data)
    elif fmt in ("csv", "tsv"):
        return read_document(data)
    else:
        return (
            f"Table extraction not supported for {path.suffix}. Supported: .pdf, .docx, .xlsx, .csv"
        )

    if not tables:
        return "No tables found in document."

    # Format output
    lines = [f"📊 Extracted {len(tables)} table(s) from {path.name}:\n"]
    for t in tables:
        page_info = f" (page {t['page']})" if t.get("page") else ""
        idx = t.get("table_index", t.get("page", "?"))
        lines.append(f"--- Table {idx}{page_info} ---")

        if t.get("note"):
            lines.append(f"  Note: {t['note']}")

        headers = t.get("headers", [])
        rows = t.get("rows", [])

        if headers:
            lines.append(" | ".join(str(h)[:30] for h in headers))
            lines.append("-" * 60)

        for row in rows[:50]:
            lines.append(" | ".join(str(c)[:30] for c in row))

        if len(rows) > 50:
            lines.append(f"... and {len(rows) - 50} more rows")
        lines.append("")

    # Also provide JSON for programmatic use
    lines.append("\n[JSON data available via tool output for further processing]")

    return "\n".join(lines)


def ocr_image(data: dict) -> str:
    """Extract text from image files (PNG, JPG, TIFF) via OCR."""
    filepath = data.get("filepath", "").strip()
    if not filepath:
        return "Please provide an image file path."

    path = Path(filepath).expanduser()
    if not path.exists():
        return f"File not found: {filepath}"

    fmt = _detect_format(path)
    if fmt != "image":
        return f"Not an image file: {path.suffix}. Supported: .png, .jpg, .jpeg, .tiff, .bmp, .gif, .webp"

    result = _ocr_image(path)
    size_kb = path.stat().st_size / 1024
    return f"📷 OCR: {path.name} ({size_kb:.1f} KB)\n\n{result}"


# === Tool Definitions ===

TOOLS = [
    {
        "name": "read_document",
        "description": (
            "Extract text from PDF, DOCX, XLSX, CSV, TXT, or image files with automatic format detection. "
            "Use this as the primary tool for reading document contents into the conversation. "
            "Supported formats are detected by file extension; use the format parameter to override. "
            "Returns plain text (truncated at 50K chars) or an error if the required library is unavailable."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "filepath": {"type": "string", "description": "Path to the file to read"},
                "format": {
                    "type": "string",
                    "description": "Override auto-detection (pdf, docx, xlsx, csv, txt, json, image)",
                },
            },
            "required": ["filepath"],
        },
        "handler": read_document,
        "category": "filereader",
    },
    {
        "name": "read_pdf",
        "description": (
            "Extract text and tables from a PDF file with page-level control. "
            "Use this instead of read_document when you need specific pages (e.g., '1-5' or '1,3,5'). "
            "Page numbers are 1-indexed; omit the pages parameter to extract all pages. "
            "Returns page-delimited text content truncated at 50K characters."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "filepath": {"type": "string", "description": "Path to PDF file"},
                "pages": {
                    "type": "string",
                    "description": "Page range: '1-5', '1,3,5', or omit for all pages (1-indexed)",
                },
            },
            "required": ["filepath"],
        },
        "handler": read_pdf,
        "category": "filereader",
    },
    {
        "name": "read_spreadsheet",
        "description": (
            "Read an XLSX or CSV spreadsheet with optional sheet selection and cell range extraction. "
            "Use this instead of read_document when you need specific sheets or cell ranges from a workbook. "
            "Supports sheet name selection, ranges like 'A1:D10', and listing all sheet names. "
            "Returns a pipe-delimited table with headers, capped at 200 rows."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "filepath": {"type": "string", "description": "Path to XLSX/CSV file"},
                "sheet": {
                    "type": "string",
                    "description": "Sheet name (for XLSX, default: active sheet)",
                },
                "range": {"type": "string", "description": "Cell range like 'A1:D10' (XLSX only)"},
                "list_sheets": {
                    "type": "boolean",
                    "default": False,
                    "description": "Also list all sheet names",
                },
            },
            "required": ["filepath"],
        },
        "handler": read_spreadsheet,
        "category": "filereader",
    },
    {
        "name": "summarize_file",
        "description": (
            "Read a file and return its extracted text prefixed with a summarization prompt. "
            "Use this when the user asks for a summary of a document rather than the raw content. "
            "Auto-detects format (PDF, DOCX, XLSX, CSV, TXT) and truncates to max_chars (default 12K). "
            "Returns file metadata followed by the extracted text ready for LLM summarization."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "filepath": {"type": "string", "description": "Path to file to summarize"},
                "max_chars": {
                    "type": "integer",
                    "default": 12000,
                    "description": "Max characters to extract",
                },
            },
            "required": ["filepath"],
        },
        "handler": summarize_file,
        "category": "filereader",
    },
    {
        "name": "extract_tables",
        "description": (
            "Extract tabular data from PDF or DOCX files into a structured, readable format. "
            "Use this when the user needs table data from documents rather than running text. "
            "Supports PDF (via PyMuPDF table detection) and DOCX (via python-docx). "
            "Returns formatted tables with headers and rows, capped at 50 rows per table."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "filepath": {"type": "string", "description": "Path to PDF or DOCX file"}
            },
            "required": ["filepath"],
        },
        "handler": extract_tables,
        "category": "filereader",
    },
    {
        "name": "ocr_image",
        "description": (
            "Extract text from image files (PNG, JPG, TIFF, BMP, GIF, WebP) via Tesseract OCR. "
            "Use this for scanned documents or photos of text that read_document cannot parse. "
            "Requires the Tesseract system package and pytesseract/Pillow Python libraries. "
            "Returns extracted text prefixed with file metadata, truncated at 50K characters."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "filepath": {
                    "type": "string",
                    "description": "Path to image file (.png, .jpg, .tiff, etc.)",
                }
            },
            "required": ["filepath"],
        },
        "handler": ocr_image,
        "category": "filereader",
    },
]
