app_kit/demo_pack.py

from __future__ import annotations

import hashlib
import json
import os
import re
import shutil
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from .storage import connect, init_db, reset_db, utc_now, copy_into_artifacts
from .index import CombinedSearchIndex, SearchHit, tokenize

HEADING_RE = re.compile(r"^(#{1,6})\s+(.+?)\s*$")
CODEBLOCK_RE = re.compile(r"```.*?```", re.S)


@dataclass
class SectionRecord:
    title: str
    slug: str
    order: int
    content: str
    page_hint: str | None
    citation: str


def slugify(text: str) -> str:
    text = text.lower().strip()
    text = re.sub(r"[^a-z0-9]+", "_", text)
    text = re.sub(r"_+", "_", text).strip("_")
    return text or "section"


def normalize_text(text: str) -> str:
    text = text.replace("\r\n", "\n")
    text = CODEBLOCK_RE.sub("", text)
    text = re.sub(r"\n{3,}", "\n\n", text)
    return text.strip()


def split_markdown_sections(text: str, default_title: str = "Overview") -> list[SectionRecord]:
    text = normalize_text(text)
    lines = text.splitlines()
    sections: list[SectionRecord] = []
    current_title = default_title
    current_lines: list[str] = []
    current_level = 0
    order = 0

    def flush() -> None:
        nonlocal order, current_lines, current_title
        content = "\n".join(current_lines).strip()
        if content:
            order += 1
            sections.append(
                SectionRecord(
                    title=current_title,
                    slug=slugify(current_title),
                    order=order,
                    content=content,
                    page_hint=None,
                    citation=f"{current_title}",
                )
            )
        current_lines = []

    for line in lines:
        match = HEADING_RE.match(line)
        if match:
            level = len(match.group(1))
            title = match.group(2).strip()
            if current_lines:
                flush()
            current_title = title
            current_level = level
            continue
        if not current_lines and not line.strip():
            continue
        current_lines.append(line)

    if current_lines:
        flush()

    if not sections:
        return [
            SectionRecord(
                title=default_title,
                slug=slugify(default_title),
                order=1,
                content=text,
                page_hint=None,
                citation=default_title,
            )
        ]
    return sections


def chunk_long_section(section: SectionRecord, max_chars: int = 1400) -> list[SectionRecord]:
    if len(section.content) <= max_chars:
        return [section]
    chunks = []
    paragraphs = [p.strip() for p in section.content.split("\n\n") if p.strip()]
    buffer: list[str] = []
    count = 0
    idx = 1
    for para in paragraphs:
        if count + len(para) > max_chars and buffer:
            content = "\n\n".join(buffer)
            chunks.append(
                SectionRecord(
                    title=f"{section.title} ({idx})",
                    slug=f"{section.slug}_{idx}",
                    order=section.order * 100 + idx,
                    content=content,
                    page_hint=section.page_hint,
                    citation=section.citation,
                )
            )
            idx += 1
            buffer = [para]
            count = len(para)
        else:
            buffer.append(para)
            count += len(para)
    if buffer:
        chunks.append(
            SectionRecord(
                title=f"{section.title} ({idx})" if idx > 1 else section.title,
                slug=f"{section.slug}_{idx}" if idx > 1 else section.slug,
                order=section.order * 100 + idx,
                content="\n\n".join(buffer),
                page_hint=section.page_hint,
                citation=section.citation,
            )
        )
    return chunks


def extract_photo_caption(path: str | Path) -> str:
    path = Path(path)
    name = path.stem.replace("_", " ").replace("-", " ")
    try:
        from PIL import Image

        with Image.open(path) as im:
            rgb = im.convert("RGB")
            sample = rgb.resize((1, 1))
            r, g, b = sample.getpixel((0, 0))
            avg = f"rgb({r},{g},{b})"
            return f"Synthetic photo '{name}' ({im.width}x{im.height}, average {avg})"
    except Exception:
        return f"Synthetic photo '{name}'"


def hash_file(path: str | Path) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(65536), b""):
            h.update(chunk)
    return h.hexdigest()


def load_manifest(pack_dir: str | Path) -> dict[str, Any]:
    pack_dir = Path(pack_dir)
    with open(pack_dir / "manifest.json", "r", encoding="utf-8") as f:
        return json.load(f)


def _resolve_uploaded_manual_path(uploaded: Any) -> Path:
    if isinstance(uploaded, Path):
        return uploaded
    if isinstance(uploaded, str):
        return Path(uploaded)
    for attr in ("path", "filepath"):
        value = getattr(uploaded, attr, None)
        if isinstance(value, str) and value:
            return Path(value)
    if isinstance(uploaded, dict):
        for key in ("path", "filepath", "name"):
            value = uploaded.get(key)
            if isinstance(value, str) and value:
                return Path(value)
    name = getattr(uploaded, "name", None)
    if isinstance(name, str) and name:
        candidate = Path(name)
        if candidate.exists():
            return candidate
    raise ValueError(f"Unsupported uploaded manual payload: {type(uploaded).__name__}")


def _read_uploaded_manual_body(path: Path) -> str:
    if path.suffix.lower() == ".pdf":
        try:
            from pypdf import PdfReader  # type: ignore[import-not-found]
        except Exception as exc:  # pragma: no cover - dependency issue is surfaced to the UI
            raise ValueError("PDF uploads require the optional 'pypdf' dependency.") from exc

        reader = PdfReader(str(path))
        pages: list[str] = []
        for page in reader.pages:
            text = page.extract_text() or ""
            if text.strip():
                pages.append(text)
        body_text = "\n\n".join(pages).strip()
        if not body_text:
            raise ValueError(f"Could not extract text from PDF manual {path.name}.")
        return body_text

    return path.read_text(encoding="utf-8", errors="ignore")


def store_uploaded_manual(conn, uploaded: Any) -> int:
    path = _resolve_uploaded_manual_path(uploaded)
    body_text = _read_uploaded_manual_body(path)
    now = utc_now()
    cur = conn.execute(
        """
        INSERT OR REPLACE INTO manuals(title, source_url, license_name, attribution, created_at, is_demo, body_text)
        VALUES (?, ?, ?, ?, ?, 0, ?)
        """,
        (
            path.stem,
            path.name,
            "uploaded",
            path.name,
            now,
            body_text,
        ),
    )
    manual_id = cur.lastrowid
    if manual_id is None:
        row = conn.execute(
            "SELECT id FROM manuals WHERE title=? AND source_url=?",
            (path.stem, path.name),
        ).fetchone()
        manual_id = int(row[0]) if row else None
    if manual_id is None:
        raise RuntimeError(f"Could not determine database id for uploaded manual {path.name}.")

    sections = split_markdown_sections(body_text, default_title=path.stem)
    final_sections: list[SectionRecord] = []
    for section in sections:
        final_sections.extend(chunk_long_section(section))
    conn.execute("DELETE FROM manual_sections WHERE manual_id = ?", (manual_id,))
    for section in final_sections:
        conn.execute(
            """
            INSERT INTO manual_sections(manual_id, section_slug, section_title, section_order, page_hint, content, citation, is_demo)
            VALUES (?, ?, ?, ?, ?, ?, ?, 0)
            """,
            (
                manual_id,
                section.slug,
                section.title,
                section.order,
                section.page_hint,
                section.content,
                f"{path.stem} :: {section.title}",
            ),
        )
    return int(manual_id)


def store_manual(conn, manual: dict[str, Any], pack_dir: Path) -> int:
    body_path = pack_dir / manual["path"]
    body_text = body_path.read_text(encoding="utf-8")
    now = utc_now()
    cur = conn.execute(
        """
        INSERT OR REPLACE INTO manuals(title, source_url, license_name, attribution, created_at, is_demo, body_text)
        VALUES (?, ?, ?, ?, ?, 1, ?)
        """,
        (
            manual["title"],
            manual["source_url"],
            manual["license_name"],
            manual["attribution"],
            now,
            body_text,
        ),
    )
    manual_id = cur.lastrowid
    if manual_id is None:
        row = conn.execute(
            "SELECT id FROM manuals WHERE title=? AND source_url=?",
            (manual["title"], manual["source_url"]),
        ).fetchone()
        manual_id = int(row[0])
    conn.execute("DELETE FROM manual_sections WHERE manual_id = ?", (manual_id,))
    sections = split_markdown_sections(body_text, default_title=manual["title"])
    final_sections: list[SectionRecord] = []
    for section in sections:
        final_sections.extend(chunk_long_section(section))
    for section in final_sections:
        conn.execute(
            """
            INSERT INTO manual_sections(manual_id, section_slug, section_title, section_order, page_hint, content, citation, is_demo)
            VALUES (?, ?, ?, ?, ?, ?, ?, 1)
            """,
            (
                manual_id,
                section.slug,
                section.title,
                section.order,
                section.page_hint,
                section.content,
                f"{manual['title']} :: {section.title}",
            ),
        )
    return int(manual_id)


def store_job(conn, job: dict[str, Any], pack_dir: Path | None = None, is_demo: bool = True) -> int:
    photo_path = None
    photo_caption = None
    if job.get("photo"):
        source = Path(pack_dir / job["photo"]) if pack_dir else Path(job["photo"])
        if source.exists() and pack_dir is None:
            photo_path = copy_into_artifacts(source)
        else:
            photo_path = str(source)
        photo_caption = extract_photo_caption(photo_path)
    response = {
        "recommendation": job.get("resolution", "insufficient evidence"),
        "expected_section_ids": job.get("expected_section_ids", []),
        "expected_section_titles": job.get("expected_section_titles", []),
        "pack_job_id": job.get("job_id"),
    }
    cur = conn.execute(
        """
        INSERT INTO jobs(
            created_at, job_title, technician, location, equipment_type, severity,
            symptom, notes, photo_path, photo_caption, response_json,
            resolution_status, linked_section_ids, linked_job_ids, is_demo
        ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        """,
        (
            utc_now(),
            job["title"],
            job.get("technician", "Demo Operator"),
            job.get("location", "Remote site"),
            job.get("equipment_type", "unknown"),
            job.get("severity", "medium"),
            job["symptom"],
            job.get("notes", ""),
            photo_path,
            photo_caption,
            json.dumps(response),
            job.get("resolution_status", "needs review"),
            json.dumps(job.get("expected_section_ids", [])),
            json.dumps(job.get("linked_job_ids", [])),
            1 if is_demo else 0,
        ),
    )
    return int(cur.lastrowid)


def ingest_demo_pack(pack_dir: str | Path, db_path: str | Path | None = None, reset: bool = True) -> dict[str, Any]:
    pack_dir = Path(pack_dir)
    manifest = load_manifest(pack_dir)
    init_db(db_path)
    if reset:
        reset_db(db_path)
        init_db(db_path)
    manual_ids: list[int] = []
    job_ids: list[int] = []
    photo_files: list[str] = []
    with connect(db_path) as conn:
        for manual in manifest.get("manuals", []):
            manual_ids.append(store_manual(conn, manual, pack_dir))
        for job in manifest.get("jobs", []):
            job_ids.append(store_job(conn, job, pack_dir=pack_dir, is_demo=True))
            if job.get("photo"):
                photo_files.append(str(pack_dir / job["photo"]))
        conn.commit()
    return {
        "pack_name": manifest.get("name", pack_dir.name),
        "manual_count": len(manual_ids),
        "job_count": len(job_ids),
        "photo_count": len(photo_files),
        "manual_ids": manual_ids,
        "job_ids": job_ids,
    }


def list_demo_packs(root: str | Path) -> list[Path]:
    root = Path(root)
    if not root.exists():
        return []
    packs = []
    for child in sorted(root.iterdir()):
        if child.is_dir() and (child / "manifest.json").exists():
            packs.append(child)
    return packs


def load_index(db_path: str | Path | None = None) -> CombinedSearchIndex:
    with connect(db_path) as conn:
        section_docs = []
        for row in conn.execute(
            """
            SELECT ms.id AS record_id, ms.section_title AS title, ms.content AS text,
                   ms.citation AS citation, ms.section_slug AS slug, m.title AS manual_title,
                   m.source_url AS source_url
            FROM manual_sections ms
            JOIN manuals m ON m.id = ms.manual_id
            ORDER BY ms.section_order, ms.id
            """
        ):
            section_docs.append(
                {
                    "record_id": row["record_id"],
                    "title": f"{row['manual_title']} — {row['title']}",
                    "text": f"{row['title']}\n\n{row['text']}",
                    "citation": row["citation"],
                    "metadata": {"slug": row["slug"], "source_url": row["source_url"]},
                }
            )
        job_docs = []
        for row in conn.execute(
            """
            SELECT id, job_title, symptom, equipment_type, location, resolution_status, photo_caption
            FROM jobs
            ORDER BY id
            """
        ):
            text = "\n".join(
                [
                    str(row["job_title"]),
                    str(row["equipment_type"]),
                    str(row["location"]),
                    str(row["symptom"]),
                    str(row["resolution_status"]),
                    str(row["photo_caption"] or ""),
                ]
            )
            job_docs.append(
                {
                    "record_id": row["id"],
                    "title": row["job_title"],
                    "text": text,
                    "citation": f"Job #{row['id']}",
                    "metadata": {"equipment_type": row["equipment_type"], "location": row["location"]},
                }
            )
    return CombinedSearchIndex(section_docs=section_docs, job_docs=job_docs)


def ensure_sample_image(path: str | Path, label: str, fill: tuple[int, int, int]) -> str:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    try:
        from PIL import Image, ImageDraw, ImageFont

        img = Image.new("RGB", (960, 540), fill)
        draw = ImageDraw.Draw(img)
        draw.rounded_rectangle((40, 40, 920, 500), radius=24, outline=(255, 255, 255), width=6)
        draw.text((70, 80), label, fill=(255, 255, 255))
        draw.text((70, 140), "Synthetic demo photo", fill=(255, 255, 255))
        img.save(path)
    except Exception:
        path.write_text(label, encoding="utf-8")
    return str(path)


def build_default_pack(root: str | Path) -> dict[str, Any]:
    """Create the demo pack if it is missing.

    This is used for bootstrapping fresh workspaces.
    """
    root = Path(root)
    root.mkdir(parents=True, exist_ok=True)
    return {"root": str(root)}