app.py

import os
import io
import json
import base64
import shutil
from datetime import datetime

from flask import Flask, request, jsonify, send_file, render_template
from flask_cors import CORS
import fitz  # PyMuPDF
from PIL import Image
from docx import Document
from docx.shared import Pt, RGBColor
from docx.enum.text import WD_ALIGN_PARAGRAPH

# ── Flask init ────────────────────────────────────────────────────────────────
app = Flask(__name__)
CORS(app)

# ── Absolute data paths ───────────────────────────────────────────────────────
UPLOAD_FOLDER = '/data/uploads'
OUTPUT_FOLDER = '/data/outputs'
META_FOLDER   = '/data/meta'

for _dir in (UPLOAD_FOLDER, OUTPUT_FOLDER, META_FOLDER):
    os.makedirs(_dir, exist_ok=True)

APP_VERSION = 'v4.2'


# ── Meta helpers ──────────────────────────────────────────────────────────────

def meta_path(filename: str) -> str:
    return os.path.join(META_FOLDER, os.path.splitext(filename)[0] + '.json')


def load_meta(filename: str) -> dict:
    path = meta_path(filename)
    if os.path.exists(path):
        with open(path, 'r', encoding='utf-8') as f:
            return json.load(f)
    return {'pages': {}}


def save_meta(filename: str, meta: dict) -> None:
    with open(meta_path(filename), 'w', encoding='utf-8') as f:
        json.dump(meta, f, ensure_ascii=False, indent=2)


# ── Text-layer detection ──────────────────────────────────────────────────────

def pdf_has_text_layer(pdf_path: str) -> bool:
    """Return True if ANY page in the PDF already contains selectable text."""
    try:
        doc = fitz.open(pdf_path)
        for page in doc:
            if page.get_text('text').strip():
                doc.close()
                return True
        doc.close()
    except Exception:
        pass
    return False


# ── APA-style header helpers ──────────────────────────────────────────────────

def apa_header(filename: str) -> str:
    stem = os.path.splitext(filename)[0]
    date = datetime.now().strftime('%Y-%m-%d')
    return (
        f'Source: Author (n.d.). {stem}. '
        f'Transcriptum Processed Data ({date}).\n'
        + '─' * 60 + '\n\n'
    )


# ── Routes ────────────────────────────────────────────────────────────────────

@app.route('/')
def index():
    return render_template('index.html', version=APP_VERSION)


# -- Version -------------------------------------------------------------------

@app.route('/api/version')
def version():
    return jsonify({'version': APP_VERSION})


# -- File list -----------------------------------------------------------------

@app.route('/api/files')
def list_files():
    """
    Scan /data/uploads on every call (no in-memory cache) so a page refresh
    always reflects the true state of the filesystem.  For each PDF found,
    cross-reference /data/outputs and /data/meta to return full linkage data.
    """
    # Build a set of output stems for fast lookup
    output_stems: set = set()
    if os.path.isdir(OUTPUT_FOLDER):
        for out_f in os.listdir(OUTPUT_FOLDER):
            output_stems.add(os.path.splitext(out_f)[0])

    files = []
    for name in sorted(os.listdir(UPLOAD_FOLDER)):
        if not name.lower().endswith('.pdf'):
            continue

        stem  = os.path.splitext(name)[0]
        meta  = load_meta(name)   # returns {} gracefully if missing
        pages = 0
        try:
            doc   = fitz.open(os.path.join(UPLOAD_FOLDER, name))
            pages = len(doc)
            doc.close()
        except Exception:
            pass

        transcribed = sum(1 for v in meta.get('pages', {}).values() if v.strip())
        has_text    = meta.get('has_text_layer', False)
        auto_output = meta.get('auto_output', '')

        # Resolve all known outputs for this source file
        known_outputs = list(meta.get('outputs', []))
        if auto_output and auto_output not in known_outputs:
            known_outputs.append(auto_output)
        # Also pick up legacy _searchable.pdf if it exists
        legacy = stem + '_searchable.pdf'
        if legacy in output_stems and legacy not in known_outputs:
            known_outputs.append(legacy)

        # Verify each output actually exists on disk
        valid_outputs = [
            o for o in known_outputs
            if os.path.exists(os.path.join(OUTPUT_FOLDER, o))
        ]

        has_meta = os.path.exists(meta_path(name))

        files.append({
            'name':        name,
            'pages':       pages,
            'transcribed': transcribed,
            'has_text':    has_text,
            'outputs':     valid_outputs,   # list of ready-to-download filenames
            'has_meta':    has_meta,
        })
    return jsonify(files)


# -- Upload --------------------------------------------------------------------

@app.route('/api/upload', methods=['POST'])
def upload_file():
    if 'file' not in request.files:
        return jsonify({'error': 'No file provided'}), 400

    f    = request.files['file']
    name = f.filename
    ext  = os.path.splitext(name)[1].lower()

    if ext == '.pdf':
        dest = os.path.join(UPLOAD_FOLDER, name)
        f.save(dest)
        doc   = fitz.open(dest)
        pages = len(doc)
        doc.close()

        # ── Smart OCR / text-layer detection ──────────────────────────────
        # If the PDF already carries a text layer (native or Transcriptum-
        # injected render_mode=3), extract its content page-by-page and
        # pre-populate the meta JSON so the transcription editor is
        # immediately filled when the user opens the document.
        if pdf_has_text_layer(dest):
            stem      = os.path.splitext(name)[0]
            date_str  = datetime.now().strftime('%Y-%m-%d_%H%M')
            out_name  = f'{stem}_{date_str}.pdf'
            out_path  = os.path.join(OUTPUT_FOLDER, out_name)
            shutil.copy2(dest, out_path)

            # Extract text from every page and store in meta.
            # page.get_text('text') retrieves both visible and invisible text;
            # for scanned + Transcriptum-injected PDFs this is the transcription.
            m = load_meta(name)
            m['has_text_layer'] = True
            m['auto_output']    = out_name
            m.setdefault('pages', {})

            doc_read = fitz.open(dest)
            for page_idx in range(len(doc_read)):
                page_text = doc_read[page_idx].get_text('text').strip()
                if page_text:
                    # Only write if no existing manual transcription for this
                    # page — preserves user edits on re-upload.
                    if str(page_idx) not in m['pages']:
                        m['pages'][str(page_idx)] = page_text
            doc_read.close()

            save_meta(name, m)

            extracted = sum(1 for v in m['pages'].values() if v.strip())
            return jsonify({
                'name':           name,
                'pages':          pages,
                'has_text_layer': True,
                'auto_output':    out_name,
                'extracted':      extracted,
                'message':        f'Text layer detected — {extracted} page(s) extracted into editor.',
            })

        return jsonify({'name': name, 'pages': pages, 'has_text_layer': False})

    if ext in ('.jpg', '.jpeg', '.png', '.tiff', '.tif'):
        img = Image.open(f.stream)
        if img.mode in ('RGBA', 'P'):
            img = img.convert('RGB')
        pdf_name = os.path.splitext(name)[0] + '.pdf'
        dest     = os.path.join(UPLOAD_FOLDER, pdf_name)
        img.save(dest, 'PDF', resolution=150)
        return jsonify({'name': pdf_name, 'pages': 1, 'has_text_layer': False})

    return jsonify({'error': 'Unsupported file type'}), 400


# -- Page render ---------------------------------------------------------------

@app.route('/api/page/<filename>/<int:page_num>')
def get_page(filename, page_num):
    pdf_path = os.path.join(UPLOAD_FOLDER, filename)
    if not os.path.exists(pdf_path):
        return jsonify({'error': 'File not found'}), 404

    doc = fitz.open(pdf_path)
    if not (0 <= page_num < len(doc)):
        doc.close()
        return jsonify({'error': 'Invalid page number'}), 400

    page = doc[page_num]
    # 2.5× gives ~180 dpi display quality; preserves original page.rect exactly
    pix  = page.get_pixmap(matrix=fitz.Matrix(2.5, 2.5), alpha=False)

    img_b64 = base64.b64encode(pix.tobytes('png')).decode()
    orig_w  = page.rect.width
    orig_h  = page.rect.height
    doc.close()

    meta      = load_meta(filename)
    page_text = meta.get('pages', {}).get(str(page_num), '')

    return jsonify({
        'image':    img_b64,
        'text':     page_text,
        'page':     page_num,
        'orig_w':   orig_w,
        'orig_h':   orig_h,
        'render_w': pix.width,
        'render_h': pix.height,
    })


# -- Save / load per-page text -------------------------------------------------

@app.route('/api/text/<filename>/<int:page_num>', methods=['GET'])
def get_text(filename, page_num):
    meta = load_meta(filename)
    return jsonify({'text': meta.get('pages', {}).get(str(page_num), ''), 'page': page_num})


@app.route('/api/text/<filename>/<int:page_num>', methods=['POST'])
def save_text(filename, page_num):
    text = (request.json or {}).get('text', '')
    meta = load_meta(filename)
    meta.setdefault('pages', {})
    if text.strip():
        meta['pages'][str(page_num)] = text
    else:
        meta['pages'].pop(str(page_num), None)
    save_meta(filename, meta)
    return jsonify({'ok': True})


# -- Internal search -----------------------------------------------------------

@app.route('/api/search/<filename>')
def search(filename):
    query = request.args.get('q', '').strip().lower()
    if not query:
        return jsonify({'results': []})

    results = []
    for page_str, text in load_meta(filename).get('pages', {}).items():
        lower = text.lower()
        idx   = 0
        while True:
            pos = lower.find(query, idx)
            if pos == -1:
                break
            start = max(0, pos - 40)
            end   = min(len(text), pos + len(query) + 40)
            results.append({'page': int(page_str), 'snippet': text[start:end], 'pos': pos})
            idx = pos + 1

    results.sort(key=lambda r: (r['page'], r['pos']))
    return jsonify({'results': results, 'query': query})


# -- Inject invisible text layer (render_mode=3) -------------------------------
#
#   Uses the EXACT rect of each page (preserves all MediaBox sizes).
#   Output filename: OriginalName_YYYY-MM-DD_HHMM.pdf  (Point 4)
# -----------------------------------------------------------------------------

@app.route('/api/inject/<filename>', methods=['POST'])
def inject(filename):
    pdf_path = os.path.join(UPLOAD_FOLDER, filename)
    if not os.path.exists(pdf_path):
        return jsonify({'error': 'File not found'}), 404

    pages_text = load_meta(filename).get('pages', {})
    if not any(v.strip() for v in pages_text.values()):
        return jsonify({'error': 'No transcription data to inject'}), 400

    doc  = fitz.open(pdf_path)
    font = fitz.Font('helv')

    for idx in range(len(doc)):
        text = pages_text.get(str(idx), '').strip()
        if not text:
            continue

        page      = doc[idx]
        page_rect = page.rect          # exact per-page dimensions
        fontsize  = 11
        line_h    = fontsize * 1.4

        tw    = fitz.TextWriter(page_rect)
        y     = fontsize + 4
        max_y = page_rect.height - 4

        for line in text.split('\n'):
            if y > max_y:
                break
            if line.strip():
                tw.append((4, y), line, font=font, fontsize=fontsize)
            y += line_h

        # render_mode=3 → invisible, fully searchable (PDF spec §9.3.6)
        tw.write_text(page, render_mode=3)

    # ── Naming: OriginalName_YYYY-MM-DD_HHMM.pdf ──────────────────────────
    stem      = os.path.splitext(filename)[0]
    timestamp = datetime.now().strftime('%Y-%m-%d_%H%M')
    out_name  = f'{stem}_{timestamp}.pdf'
    out_path  = os.path.join(OUTPUT_FOLDER, out_name)
    doc.save(out_path, garbage=4, deflate=True)
    doc.close()

    # Record output reference in meta for deep-cleanup
    m = load_meta(filename)
    m.setdefault('outputs', [])
    if out_name not in m['outputs']:
        m['outputs'].append(out_name)
    save_meta(filename, m)

    return jsonify({'ok': True, 'output': out_name})



# -- Inject single page (active page only → single-page PDF output) -----------

@app.route('/api/inject_page/<filename>/<int:page_num>', methods=['POST'])
def inject_page(filename, page_num):
    """
    Inject invisible text into ONE page and export it as a standalone
    single-page PDF.  The source document is not modified.
    """
    pdf_path = os.path.join(UPLOAD_FOLDER, filename)
    if not os.path.exists(pdf_path):
        return jsonify({'error': 'File not found'}), 404

    pages_text = load_meta(filename).get('pages', {})
    page_text  = pages_text.get(str(page_num), '').strip()
    if not page_text:
        return jsonify({'error': f'No transcription text for page {page_num + 1}'}), 400

    src_doc = fitz.open(pdf_path)
    if page_num >= len(src_doc):
        src_doc.close()
        return jsonify({'error': 'Invalid page number'}), 400

    # Build a new single-page document from the source page
    out_doc = fitz.open()
    out_doc.insert_pdf(src_doc, from_page=page_num, to_page=page_num)
    src_doc.close()

    page      = out_doc[0]
    page_rect = page.rect
    font      = fitz.Font('helv')
    fontsize  = 11
    line_h    = fontsize * 1.4
    tw        = fitz.TextWriter(page_rect)
    y         = fontsize + 4
    max_y     = page_rect.height - 4

    for line in page_text.split('\n'):
        if y > max_y:
            break
        if line.strip():
            tw.append((4, y), line, font=font, fontsize=fontsize)
        y += line_h

    tw.write_text(page, render_mode=3)

    stem      = os.path.splitext(filename)[0]
    timestamp = datetime.now().strftime('%Y-%m-%d_%H%M')
    out_name  = f'{stem}_p{page_num + 1}_{timestamp}.pdf'
    out_path  = os.path.join(OUTPUT_FOLDER, out_name)
    out_doc.save(out_path, garbage=4, deflate=True)
    out_doc.close()

    return jsonify({'ok': True, 'output': out_name})

# -- Download ------------------------------------------------------------------

@app.route('/api/download/<filename>')
def download(filename):
    for folder in (OUTPUT_FOLDER, UPLOAD_FOLDER):
        path = os.path.join(folder, filename)
        if os.path.exists(path):
            return send_file(path, as_attachment=True, download_name=filename)
    return jsonify({'error': 'File not found'}), 404


# -- Export .txt ---------------------------------------------------------------

@app.route('/api/export/txt/<filename>')
def export_txt(filename):
    pages = load_meta(filename).get('pages', {})
    if not pages:
        return jsonify({'error': 'No transcription data'}), 400

    stem      = os.path.splitext(filename)[0]
    timestamp = datetime.now().strftime('%Y-%m-%d_%H%M')
    lines     = [apa_header(filename)]

    max_page = max(int(k) for k in pages)
    for i in range(max_page + 1):
        text = pages.get(str(i), '').strip()
        if text:
            lines += [f'--- Page {i + 1} ---', text, '']

    buf = io.BytesIO('\n'.join(lines).encode('utf-8'))
    buf.seek(0)
    return send_file(buf, as_attachment=True,
                     download_name=f'{stem}_{timestamp}_transcription.txt',
                     mimetype='text/plain')


# -- Export .docx --------------------------------------------------------------

@app.route('/api/export/docx/<filename>')
def export_docx(filename):
    pages = load_meta(filename).get('pages', {})
    if not pages:
        return jsonify({'error': 'No transcription data'}), 400

    stem      = os.path.splitext(filename)[0]
    timestamp = datetime.now().strftime('%Y-%m-%d_%H%M')
    date_str  = datetime.now().strftime('%Y-%m-%d')

    doc_out = Document()
    style   = doc_out.styles['Normal']
    style.font.name = 'Garamond'
    style.font.size = Pt(11)

    # APA-style header — matches the .txt export header exactly
    apa_p = doc_out.add_paragraph()
    apa_r = apa_p.add_run(f'Source: Author (n.d.). {stem}. Transcriptum Processed Data ({date_str}).')
    apa_r.font.size      = Pt(9)
    apa_r.font.color.rgb = RGBColor(0x80, 0x70, 0x60)
    apa_r.italic         = True
    # Separator line to match .txt format
    sep_p = doc_out.add_paragraph('─' * 60)
    sep_p.runs[0].font.size      = Pt(8)
    sep_p.runs[0].font.color.rgb = RGBColor(0x80, 0x70, 0x60)
    doc_out.add_paragraph()

    max_page = max(int(k) for k in pages)
    for i in range(max_page + 1):
        text = pages.get(str(i), '').strip()
        if not text:
            continue
        h = doc_out.add_heading(f'Page {i + 1}', level=1)
        h.runs[0].font.size = Pt(13)
        doc_out.add_paragraph(text)
        doc_out.add_paragraph()

    buf = io.BytesIO()
    doc_out.save(buf)
    buf.seek(0)
    return send_file(buf, as_attachment=True,
                     download_name=f'{stem}_{timestamp}_transcription.docx',
                     mimetype='application/vnd.openxmlformats-officedocument.wordprocessingml.document')


# -- Progress ------------------------------------------------------------------

@app.route('/api/progress/<filename>')
def progress(filename):
    pdf_path = os.path.join(UPLOAD_FOLDER, filename)
    if not os.path.exists(pdf_path):
        return jsonify({'error': 'Not found'}), 404
    doc   = fitz.open(pdf_path)
    total = len(doc)
    doc.close()
    meta        = load_meta(filename)
    transcribed = sum(1 for v in meta.get('pages', {}).values() if v.strip())
    return jsonify({'total': total, 'transcribed': transcribed})


# -- Delete (deep cleanup — original logic preserved) --------------------------
#
#   Removes every artefact linked to a source PDF:
#     1. /data/uploads/{filename}
#     2. /data/outputs/* — legacy _searchable, timestamped, and meta['outputs']
#     3. /data/meta/{stem}.json
# -----------------------------------------------------------------------------

@app.route('/api/delete/<filename>', methods=['DELETE'])
def delete_file(filename):
    if os.sep in filename or filename.startswith('.'):
        return jsonify({'error': 'Invalid filename'}), 400

    stem    = os.path.splitext(filename)[0]
    deleted = []
    skipped = []

    # 1 — source upload
    src = os.path.join(UPLOAD_FOLDER, filename)
    if os.path.exists(src):
        os.remove(src)
        deleted.append(src)
    else:
        skipped.append(src)

    # 2 — all linked output files
    outputs_to_delete = set()

    # A: legacy _searchable.pdf
    legacy = os.path.join(OUTPUT_FOLDER, stem + '_searchable.pdf')
    if os.path.exists(legacy):
        outputs_to_delete.add(legacy)

    # B: references stored in meta['outputs']
    stored_meta = load_meta(filename)
    for ref in stored_meta.get('outputs', []):
        candidate = os.path.join(OUTPUT_FOLDER, os.path.basename(ref))
        if os.path.exists(candidate):
            outputs_to_delete.add(candidate)

    # C: auto_output from text-layer detection
    auto = stored_meta.get('auto_output', '')
    if auto:
        candidate = os.path.join(OUTPUT_FOLDER, auto)
        if os.path.exists(candidate):
            outputs_to_delete.add(candidate)

    # D: scan outputs dir for explicit stem matches
    if os.path.isdir(OUTPUT_FOLDER):
        for out_file in os.listdir(OUTPUT_FOLDER):
            out_stem = os.path.splitext(out_file)[0]
            if out_stem == stem or out_stem == stem + '_searchable':
                outputs_to_delete.add(os.path.join(OUTPUT_FOLDER, out_file))

    for out_path in outputs_to_delete:
        if os.path.exists(out_path):
            os.remove(out_path)
            deleted.append(out_path)

    # 3 — metadata JSON
    mpath = meta_path(filename)
    if os.path.exists(mpath):
        os.remove(mpath)
        deleted.append(mpath)
    else:
        skipped.append(mpath)

    return jsonify({'ok': True, 'deleted': deleted, 'skipped': skipped})


# ── Entry point ───────────────────────────────────────────────────────────────
if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=False)