#!/usr/bin/env python3
"""
Спрашивает по документам Paperless-ngx через поиск по OCR и отвечает через Ollama.
Использование: PAPERLESS_URL=... PAPERLESS_TOKEN=... python paperless-ollama-ask.py "номер паспорта?"
"""
import json
import os
import sys
import urllib.parse
import urllib.request


def env(name: str, default: str = "") -> str:
    v = os.environ.get(name, default).strip()
    if not v and name.startswith("PAPERLESS_"):
        raise SystemExit(f"Set {name} (e.g. PAPERLESS_URL=http://192.168.1.104:8000 PAPERLESS_TOKEN=...)")
    return v


# Слова, которые редко есть в OCR тексте документов — выкидываем при запасном поиске
_STOP_WORDS = frozenset(
    "кем какой какая какие как где когда что кто чей чья чьи откуда куда зачем почему сколько".split()
)


def search_query_for_paperless(question: str) -> str:
    """Убираем пунктуацию, которая может ломать full-text поиск Paperless (?, ! и т.д.)."""
    s = question.strip().rstrip("?!.;")
    return s if s else question


def fallback_search_query(question: str, max_words: int = 4) -> str:
    """Короткий запрос для повторного поиска: ключевые слова без вопросительных слов."""
    q = search_query_for_paperless(question)
    words = [w for w in q.split() if len(w) >= 2 and w.lower() not in _STOP_WORDS]
    return " ".join(words[:max_words]) if words else q


def paperless_search(base_url: str, token: str, query: str, max_docs: int = 5) -> list[dict]:
    """Full-text поиск по документам, возвращает список с полем content (OCR)."""
    base_url = base_url.rstrip("/")
    search_q = search_query_for_paperless(query)
    query_encoded = urllib.parse.quote(search_q, encoding="utf-8", safe="")
    url = f"{base_url}/api/documents/?query={query_encoded}&page_size={max_docs}"
    req = urllib.request.Request(url)
    req.add_header("Authorization", f"Token {token}")
    req.add_header("Accept", "application/json; version=6")
    if os.environ.get("PAPERLESS_DEBUG"):
        print(f"[DEBUG] GET {url}", file=sys.stderr)
    with urllib.request.urlopen(req, timeout=30) as r:
        data = r.read().decode("utf-8")
    out = json.loads(data)
    results = out.get("results") or []
    if os.environ.get("PAPERLESS_DEBUG"):
        print(f"[DEBUG] count={out.get('count', 0)} results={len(results)}", file=sys.stderr)
    return results


def ollama_generate(ollama_url: str, model: str, prompt: str, timeout: int = 120) -> str:
    """Один запрос к Ollama /api/generate, возвращает response."""
    ollama_url = ollama_url.rstrip("/")
    url = f"{ollama_url}/api/generate"
    body = {"model": model, "prompt": prompt, "stream": False}
    data = json.dumps(body).encode()
    req = urllib.request.Request(url, data=data, method="POST")
    req.add_header("Content-Type", "application/json")
    with urllib.request.urlopen(req, timeout=timeout) as r:
        out = json.loads(r.read().decode("utf-8"))
    return (out.get("response") or "").strip()


def main():
    question = " ".join(sys.argv[1:]).strip()
    if not question:
        print("Usage: PAPERLESS_URL=... PAPERLESS_TOKEN=... python paperless-ollama-ask.py 'ваш вопрос'", file=sys.stderr)
        sys.exit(1)

    base_url = env("PAPERLESS_URL", "http://localhost:8000")
    token = env("PAPERLESS_TOKEN")
    ollama_url = env("OLLAMA_URL", "http://localhost:11434")
    model = env("OLLAMA_MODEL", "saiga")
    max_docs = int(env("PAPERLESS_MAX_DOCS", "5"))

    # Поиск по OCR: сначала полный вопрос, при пустом ответе — по ключевым словам
    try:
        docs = paperless_search(base_url, token, question, max_docs=max_docs)
        if not docs:
            fallback = fallback_search_query(question)
            if fallback and fallback != search_query_for_paperless(question):
                docs = paperless_search(base_url, token, fallback, max_docs=max_docs)
                if os.environ.get("PAPERLESS_DEBUG") and docs:
                    print(f"[DEBUG] fallback query '{fallback}' found {len(docs)} docs", file=sys.stderr)
    except urllib.error.HTTPError as e:
        print(f"Paperless API error: {e.code} {e.reason}", file=sys.stderr)
        if e.code == 401:
            print("Check PAPERLESS_TOKEN (My Profile → API token in Paperless UI).", file=sys.stderr)
        sys.exit(2)
    except Exception as e:
        print(f"Paperless request failed: {e}", file=sys.stderr)
        sys.exit(2)

    if not docs:
        print("По документам ничего не найдено.")
        return

    # Собираем текст из документов (OCR content)
    parts = []
    for i, d in enumerate(docs, 1):
        title = d.get("title") or f"Документ {d.get('id')}"
        content = (d.get("content") or "").strip()
        if not content:
            content = "(текст пустой)"
        parts.append(f"[Документ {i}: {title}]\n{content[:8000]}")  # ограничиваем длину

    docs_text = "\n\n---\n\n".join(parts)

    prompt = f"""Ниже текст из документов (OCR). Ответь на вопрос пользователя, опираясь только на этот текст.
Если в тексте нет ответа — напиши "В документах не найдено".
Отвечай кратко, по делу.

Документы:

{docs_text}

Вопрос: {question}

Ответ:"""

    try:
        answer = ollama_generate(ollama_url, model, prompt)
    except urllib.error.URLError as e:
        print(f"Ollama недоступна ({ollama_url}): {e}", file=sys.stderr)
        sys.exit(3)
    except Exception as e:
        print(f"Ollama error: {e}", file=sys.stderr)
        sys.exit(3)

    print(answer)


if __name__ == "__main__":
    main()