# -*- coding: utf-8 -*-
"""Lokalny czat z PDF-em po polsku (RAG) na Ollamie. Bez internetu, bez kluczy.
Wymaga: pip install pypdf numpy ; ollama pull bge-m3 ; ollama pull gemma4:12b-it-qat
Użycie:
  python czat_z_pdf.py dokument.pdf                    -> tryb rozmowy
  python czat_z_pdf.py dokument.pdf "Twoje pytanie"    -> jedna odpowiedź
Zmienne środowiskowe:
  MODEL  - model do odpowiedzi (domyślnie gemma4:12b-it-qat)
  EMB    - model wyszukiwania (domyślnie bge-m3)
  TOPK   - ile fragmentów dać modelowi (domyślnie 5)
  PELNY  - 1 = zamiast wyszukiwania wklej cały dokument (krótkie pliki, kontekst 65 536 tokenów)"""
import sys, os, re, json, pathlib, urllib.request
import numpy as np
from pypdf import PdfReader

OLLAMA = os.environ.get('OLLAMA_URL', 'http://127.0.0.1:11434')
MODEL = os.environ.get('MODEL', 'gemma4:12b-it-qat')
EMB = os.environ.get('EMB', 'bge-m3')
TOPK = int(os.environ.get('TOPK', '5'))
PELNY = os.environ.get('PELNY', '0') == '1'

def api(path, body):
    r = urllib.request.Request(OLLAMA + path, data=json.dumps(body).encode(), headers={'Content-Type': 'application/json'})
    return json.load(urllib.request.urlopen(r, timeout=900))

def tekst_z_pdf(plik):
    t = '\n'.join(p.extract_text() or '' for p in PdfReader(plik).pages)
    t = re.sub(r'(\w)-\n(\w)', r'\1\2', t)              # dzielenie wyrazów na końcu linii
    t = re.sub(r'[ \t]+', ' ', t)
    t = re.sub(r'\n(?!Art\. |\d+\. |\d+\) )', ' ', t)   # sklejam linie, zostawiam początki artykułów i punktów
    return t

def kawalki(t, maks=1400, zakladka=200):
    # najpierw tnę po artykułach, długie artykuły dzielę dalej z zakładką
    czesci = re.split(r'\n(?=Art\. \d+[a-z]*\. )', t)
    out = []
    for c in czesci:
        c = c.strip()
        if not c: continue
        naglowek = (re.match(r'Art\. \d+[a-z]*\.', c) or [''])[0]
        i = 0
        while i < len(c):
            kawal = c[i:i + maks]
            if i > 0 and naglowek: kawal = f'({naglowek} - ciąg dalszy) ' + kawal
            out.append(kawal)
            i += maks - zakladka
    return out

def indeks(plik):
    cache = pathlib.Path(plik).with_suffix(f'.{EMB}.npz')
    if cache.exists():
        d = np.load(cache, allow_pickle=True); return list(d['k']), d['v']
    k = kawalki(tekst_z_pdf(plik))
    v = []
    for i in range(0, len(k), 16):
        v += api('/api/embed', {'model': EMB, 'input': k[i:i + 16]})['embeddings']
    v = np.array(v, dtype=np.float32); v /= np.linalg.norm(v, axis=1, keepdims=True)
    np.savez(cache, k=np.array(k, dtype=object), v=v)
    return k, v

def szukaj(pyt, k, v, n=TOPK):
    q = np.array(api('/api/embed', {'model': EMB, 'input': [pyt]})['embeddings'][0], dtype=np.float32)
    q /= np.linalg.norm(q)
    idx = np.argsort(-(v @ q))[:n]
    return [k[i] for i in idx]

SYSTEM = ('Odpowiadasz po polsku wyłącznie na podstawie fragmentów dokumentu podanych niżej. '
          'Podaj numer artykułu, z którego bierzesz odpowiedź. Jeśli fragmenty nie zawierają odpowiedzi, '
          'napisz: "Nie znalazłem tego w dokumencie." Odpowiadaj krótko, w 1-3 zdaniach.')

def odpowiedz(pyt, k, v, model=MODEL, ctx=8192, caly=None):
    if caly is not None:
        fr, kontekst, ctx = [], f'DOKUMENT:\n{caly}', 65536
    else:
        fr = szukaj(pyt, k, v)
        kontekst = '\n\n'.join(f'[Fragment {i + 1}]\n{f}' for i, f in enumerate(fr))
    msg = [{'role': 'system', 'content': SYSTEM}, {'role': 'user', 'content': f'{kontekst}\n\nPytanie: {pyt}'}]
    body = {'model': model, 'messages': msg, 'stream': False, 'options': {'num_ctx': ctx, 'temperature': 0}}
    try: r = api('/api/chat', {**body, 'think': False})
    except Exception: r = api('/api/chat', body)
    return r['message']['content'], fr

if __name__ == '__main__':
    sys.stdout.reconfigure(encoding='utf-8')
    if len(sys.argv) < 2: print(__doc__); sys.exit(1)
    if PELNY:
        caly, k, v = tekst_z_pdf(sys.argv[1]), None, None
        print(f'Dokument: {len(caly)} znaków w całości, model {MODEL}')
    else:
        caly = None; k, v = indeks(sys.argv[1])
        print(f'Dokument: {len(k)} fragmentów, model {MODEL}, wyszukiwanie {EMB}')
    if len(sys.argv) > 2:
        print(odpowiedz(sys.argv[2], k, v, caly=caly)[0]); sys.exit()
    while True:
        try: pyt = input('\nPytanie (Enter kończy): ').strip()
        except EOFError: break
        if not pyt: break
        print(odpowiedz(pyt, k, v, caly=caly)[0])