app/documents.py

from __future__ import annotations

import csv
import html.parser
import json
import mimetypes
import re
import zipfile
from dataclasses import dataclass
from pathlib import Path
from typing import Any
from uuid import uuid4
from xml.etree import ElementTree

from app.config import ApiSettings


SUPPORTED_EXTENSIONS = {
    ".txt",
    ".md",
    ".markdown",
    ".json",
    ".csv",
    ".html",
    ".htm",
    ".pdf",
    ".docx",
}
TEXT_EXTENSIONS = {".txt", ".md", ".markdown"}
TARGET_CHARS = 2200
OVERLAP_CHARS = 200


@dataclass(frozen=True)
class ParsedChunk:
    page: int | None
    section: str | None
    row_start: int | None
    row_end: int | None
    char_start: int
    char_end: int
    text: str


@dataclass(frozen=True)
class ParsedDocument:
    title: str
    parser: str
    text: str
    chunks: list[ParsedChunk]


def infer_mime_type(filename: str, fallback: str | None = None) -> str:
    guessed, _ = mimetypes.guess_type(filename)
    return fallback or guessed or "application/octet-stream"


def validate_upload(filename: str, size_bytes: int, settings: ApiSettings) -> None:
    suffix = Path(filename).suffix.lower()
    if suffix not in SUPPORTED_EXTENSIONS:
        raise ValueError(f"Unsupported file type: {suffix or 'unknown'}")
    if size_bytes <= 0:
        raise ValueError("Uploaded file is empty.")
    if size_bytes > settings.max_upload_bytes:
        raise ValueError(
            f"Uploaded file exceeds {settings.max_upload_bytes} byte limit."
        )


def parse_document_file(
    path: Path,
    *,
    filename: str,
    mime_type: str,
    settings: ApiSettings,
) -> ParsedDocument:
    suffix = Path(filename).suffix.lower()
    if suffix in TEXT_EXTENSIONS:
        return _parse_text(path, filename, settings)
    if suffix == ".json":
        return _parse_json(path, filename, settings)
    if suffix == ".csv":
        return _parse_csv(path, filename, settings)
    if suffix in {".html", ".htm"}:
        return _parse_html(path, filename, settings)
    if suffix == ".docx":
        return _parse_docx(path, filename, settings)
    if suffix == ".pdf":
        return _parse_pdf(path, filename, settings)
    raise ValueError(f"Unsupported file type for parsing: {mime_type}")


def chunks_to_records(
    chunks: list[ParsedChunk], *, file_id: str
) -> list[dict[str, Any]]:
    records = []
    for index, chunk in enumerate(chunks):
        records.append(
            {
                "id": f"chk_{uuid4().hex}",
                "file_id": file_id,
                "chunk_index": index,
                "page": chunk.page,
                "section": chunk.section,
                "row_start": chunk.row_start,
                "row_end": chunk.row_end,
                "char_start": chunk.char_start,
                "char_end": chunk.char_end,
                "text": chunk.text,
                "token_count": _estimate_tokens(chunk.text),
                "embedding_id": None,
            }
        )
    return records


def _parse_text(path: Path, filename: str, settings: ApiSettings) -> ParsedDocument:
    text = _read_text(path, settings)
    return ParsedDocument(
        title=filename,
        parser="text",
        text=text,
        chunks=_chunk_text(text, section=_first_heading(text)),
    )


def _parse_json(path: Path, filename: str, settings: ApiSettings) -> ParsedDocument:
    raw = _read_text(path, settings)
    try:
        data = json.loads(raw)
        text = json.dumps(data, indent=2, ensure_ascii=False)
    except json.JSONDecodeError:
        text = raw
    text = _cap_text(text, settings)
    return ParsedDocument(
        title=filename,
        parser="json",
        text=text,
        chunks=_chunk_text(text, section="JSON"),
    )


def _parse_csv(path: Path, filename: str, settings: ApiSettings) -> ParsedDocument:
    raw = _read_text(path, settings)
    rows = list(csv.reader(raw.splitlines()))
    if not rows:
        raise ValueError("CSV has no rows.")

    header = rows[0]
    data_rows = rows[1:]
    summary = [
        f"CSV file: {filename}",
        f"Columns: {', '.join(header)}",
        f"Rows: {len(data_rows)}",
    ]
    chunks = [
        ParsedChunk(
            page=None,
            section="CSV summary",
            row_start=None,
            row_end=None,
            char_start=0,
            char_end=len("\n".join(summary)),
            text="\n".join(summary),
        )
    ]

    row_group_size = 50
    cursor = chunks[0].char_end
    for start in range(0, len(data_rows), row_group_size):
        group = data_rows[start : start + row_group_size]
        lines = [",".join(header)]
        lines.extend(",".join(row) for row in group)
        text = "\n".join(lines)
        chunks.append(
            ParsedChunk(
                page=None,
                section="CSV rows",
                row_start=start + 1,
                row_end=start + len(group),
                char_start=cursor,
                char_end=cursor + len(text),
                text=text,
            )
        )
        cursor += len(text)
        if len(chunks) >= settings.max_chunks_per_file:
            break

    document_text = "\n\n".join(chunk.text for chunk in chunks)
    return ParsedDocument(
        title=filename,
        parser="csv",
        text=_cap_text(document_text, settings),
        chunks=chunks[: settings.max_chunks_per_file],
    )


def _parse_html(path: Path, filename: str, settings: ApiSettings) -> ParsedDocument:
    raw = _read_text(path, settings)
    parser = _HTMLTextParser()
    parser.feed(raw)
    text = _cap_text(parser.text(), settings)
    return ParsedDocument(
        title=parser.title or filename,
        parser="html",
        text=text,
        chunks=_chunk_text(text, section=parser.title),
    )


def _parse_docx(path: Path, filename: str, settings: ApiSettings) -> ParsedDocument:
    try:
        with zipfile.ZipFile(path) as archive:
            xml = archive.read("word/document.xml")
    except Exception as exc:
        raise ValueError("Could not read DOCX document.xml.") from exc

    root = ElementTree.fromstring(xml)
    namespace = "{http://schemas.openxmlformats.org/wordprocessingml/2006/main}"
    paragraphs: list[str] = []
    for paragraph in root.iter(f"{namespace}p"):
        texts = [node.text or "" for node in paragraph.iter(f"{namespace}t")]
        line = "".join(texts).strip()
        if line:
            paragraphs.append(line)
    text = _cap_text("\n\n".join(paragraphs), settings)
    return ParsedDocument(
        title=filename,
        parser="docx",
        text=text,
        chunks=_chunk_text(text, section=_first_heading(text)),
    )


def _parse_pdf(path: Path, filename: str, settings: ApiSettings) -> ParsedDocument:
    try:
        from pypdf import PdfReader
    except ImportError as exc:
        raise ValueError("PDF parsing requires the pypdf package.") from exc

    reader = PdfReader(str(path))
    chunks: list[ParsedChunk] = []
    page_texts: list[str] = []
    cursor = 0
    for page_index, page in enumerate(reader.pages, start=1):
        text = (page.extract_text() or "").strip()
        if not text:
            continue
        page_texts.append(text)
        for chunk in _chunk_text(text, page=page_index, start_offset=cursor):
            chunks.append(chunk)
        cursor += len(text) + 2
        if len(chunks) >= settings.max_chunks_per_file:
            break
    if not chunks:
        raise ValueError("PDF text extraction returned no readable text.")
    document_text = _cap_text("\n\n".join(page_texts), settings)
    return ParsedDocument(
        title=filename,
        parser="pdf",
        text=document_text,
        chunks=chunks[: settings.max_chunks_per_file],
    )


def _read_text(path: Path, settings: ApiSettings) -> str:
    data = path.read_bytes()
    for encoding in ("utf-8", "utf-8-sig", "latin-1"):
        try:
            return _cap_text(data.decode(encoding), settings)
        except UnicodeDecodeError:
            continue
    raise ValueError("Could not decode file as text.")


def _cap_text(text: str, settings: ApiSettings) -> str:
    cleaned = text.replace("\x00", "").strip()
    return cleaned[: settings.max_parsed_chars]


def _chunk_text(
    text: str,
    *,
    page: int | None = None,
    section: str | None = None,
    start_offset: int = 0,
) -> list[ParsedChunk]:
    paragraphs = [part.strip() for part in re.split(r"\n{2,}", text) if part.strip()]
    if not paragraphs:
        return []

    chunks: list[ParsedChunk] = []
    buffer: list[str] = []
    chunk_start = start_offset
    cursor = start_offset

    def flush(end_cursor: int) -> None:
        nonlocal buffer, chunk_start
        chunk_text = "\n\n".join(buffer).strip()
        if not chunk_text:
            return
        chunks.append(
            ParsedChunk(
                page=page,
                section=section,
                row_start=None,
                row_end=None,
                char_start=chunk_start,
                char_end=end_cursor,
                text=chunk_text,
            )
        )
        overlap = chunk_text[-OVERLAP_CHARS:] if len(chunk_text) > OVERLAP_CHARS else ""
        buffer = [overlap] if overlap else []
        chunk_start = max(start_offset, end_cursor - len(overlap))

    for paragraph in paragraphs:
        paragraph_start = cursor
        paragraph_end = cursor + len(paragraph)
        candidate = "\n\n".join([*buffer, paragraph]).strip()
        if buffer and len(candidate) > TARGET_CHARS:
            flush(paragraph_start)
        if not buffer:
            chunk_start = paragraph_start
        buffer.append(paragraph)
        cursor = paragraph_end + 2

    flush(cursor)
    return chunks


def _estimate_tokens(text: str) -> int:
    return max(1, len(text) // 4)


def _first_heading(text: str) -> str | None:
    for line in text.splitlines():
        cleaned = line.strip().lstrip("#").strip()
        if cleaned:
            return cleaned[:120]
    return None


class _HTMLTextParser(html.parser.HTMLParser):
    def __init__(self) -> None:
        super().__init__()
        self._skip_depth = 0
        self._in_title = False
        self._title_parts: list[str] = []
        self._parts: list[str] = []

    @property
    def title(self) -> str | None:
        title = " ".join(" ".join(self._title_parts).split())
        return title or None

    def handle_starttag(self, tag: str, attrs) -> None:
        if tag in {"script", "style", "noscript"}:
            self._skip_depth += 1
        if tag == "title":
            self._in_title = True
        if tag in {"p", "div", "section", "article", "br", "h1", "h2", "h3", "li"}:
            self._parts.append("\n")

    def handle_endtag(self, tag: str) -> None:
        if tag in {"script", "style", "noscript"} and self._skip_depth:
            self._skip_depth -= 1
        if tag == "title":
            self._in_title = False
        if tag in {"p", "div", "section", "article", "h1", "h2", "h3", "li"}:
            self._parts.append("\n")

    def handle_data(self, data: str) -> None:
        if self._skip_depth:
            return
        cleaned = " ".join(data.split())
        if not cleaned:
            return
        if self._in_title:
            self._title_parts.append(cleaned)
        self._parts.append(cleaned)

    def text(self) -> str:
        return "\n".join(
            line.strip() for line in " ".join(self._parts).splitlines() if line.strip()
        )