src/pageparse/extract.py

from __future__ import annotations

from pageparse.config import settings
from pageparse.schema import SourceParseResult, Record


SCHEMA_DETECTION_KEYWORDS = {
    "todo": ["todo", "to-do", "task", "deadline", "due", "priority", "action item"],
    "meeting": ["meeting", "agenda", "speaker", "minutes", "attendee", "discussion"],
    "recipe": ["recipe", "ingredient", "tbsp", "tsp", "cup", "cook", "bake", "preheat"],
    "survey": ["survey", "questionnaire", "question", "response", "feedback"],
    "notes": ["note", "lecture", "chapter", "section", "topic"],
}

DEFAULT_SCHEMA_PROMPTS = {
    "meeting": (
        "Extract records of type 'action_item', 'note', or 'transcript_segment' from the following text. "
        "Include speaker, timestamp, priority, and due_date where present. "
        "Do NOT change any words. Preserve original text exactly."
    ),
    "recipe": (
        "Extract ingredients as type 'key_value' (category 'Ingredient') and "
        "steps as type 'note' (category 'Instruction'). "
        "Do NOT change any words. Preserve original text exactly."
    ),
    "survey": (
        "Extract questions as type 'note' (category 'Question') and "
        "responses as type 'key_value' (category 'Response'). "
        "Do NOT change any words. Preserve original text exactly."
    ),
    "notes": (
        "Extract key topics as type 'note' with category from the following notes. "
        "Include any action items with priority and due_date where present. "
        "Do NOT change any words. Preserve original text exactly."
    ),
}


class Extractor:
    def __init__(self) -> None:
        self._grammar: str | None = None
        grammar_path = settings.grammar_path()
        if grammar_path.exists():
            self._grammar = grammar_path.read_text(encoding="utf-8")

    def detect_schema_type(self, raw_text: str) -> str:
        if not settings.auto_schema:
            return "todo"
        text_lower = raw_text.lower()
        scores: dict[str, int] = {}
        for schema_type, keywords in SCHEMA_DETECTION_KEYWORDS.items():
            score = sum(text_lower.count(kw) for kw in keywords)
            if score > 0:
                scores[schema_type] = score
        if not scores:
            return "todo"
        return max(scores, key=scores.get)

    def extract(
        self, raw_text: str, source_file: str, schema_type: str | None = None
    ) -> SourceParseResult:
        if schema_type is None or schema_type == "auto":
            schema_type = self.detect_schema_type(raw_text)

        result = SourceParseResult(
            source_file=source_file,
            source_type="image",
            records=[],
        )

        if not raw_text.strip() or raw_text in ("[UNCLEAR]", "[inaudible]"):
            return result

        result.records = self._line_based_fallback(raw_text)

        if self._grammar and raw_text.strip():
            try:
                structured = self._ollama_structured_extract(raw_text, source_file, schema_type)
                if structured and structured.records:
                    result = structured
                    result.source_file = source_file
            except Exception:
                pass

        return result

    def _ollama_structured_extract(
        self, raw_text: str, source_file: str, schema_type: str
    ) -> SourceParseResult | None:
        import json
        import urllib.request

        user_msg = DEFAULT_SCHEMA_PROMPTS.get(
            schema_type,
            "Extract tasks with due_date, priority, status, and category from the following text. "
            "Do NOT change any words. Preserve original text exactly."
        )

        system_msg = "You are a structured data extractor. Output only valid JSON matching the schema."
        prompt = f"{system_msg}\n\n{user_msg}\n\nExtracted Text:\n{raw_text}\n\nSource File: {source_file}"

        payload = {
            "model": settings.slm_model,
            "prompt": prompt,
            "stream": False,
            "options": {"temperature": 0.1},
        }

        url = f"{settings.ollama_url}/api/generate"
        data = json.dumps(payload).encode("utf-8")
        req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"})
        with urllib.request.urlopen(req, timeout=30) as response:
            body = response.read().decode("utf-8")
            res = json.loads(body)
            raw = res.get("response", "").strip()

        if raw:
            try:
                return SourceParseResult.model_validate_json(raw)
            except Exception:
                pass
        return None

    def _line_based_fallback(self, raw_text: str) -> list[Record]:
        records: list[Record] = []
        lines = raw_text.strip().split("\n")
        for line in lines:
            line = line.strip()
            if not line:
                continue
            records.append(
                Record(
                    type="note",
                    content=line,
                    confidence=0.95,
                )
            )
        if not records:
            records.append(Record(type="note", content=raw_text.strip()))
        return records