src/pageparse/pipelines.py

from __future__ import annotations

import hashlib
import json
import shutil
import subprocess
import tempfile
import xml.etree.ElementTree as ET
import zipfile
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

import cv2
import numpy as np

from pageparse.config import settings
from pageparse.ingest import load_image
from pageparse.ocr.printed import PrintedOCR
import pageparse.preprocess as pp_module
from pageparse.preprocess import preprocess, adaptive_preprocess


def _get_uploads_dir() -> Path:
    here = Path(__file__).resolve().parent.parent.parent
    uploads_dir = here / "web" / "static" / "uploads"
    if not (here / "web" / "static").exists():
        uploads_dir = Path.cwd() / "web" / "static" / "uploads"
    return uploads_dir


def _try_handwriting_ocr(image):
    try:
        from pageparse.ocr.handwriting import HandwritingOCR
        ocr = HandwritingOCR()
        return ocr.recognize(image)
    except Exception:
        return ""


def _ocr_image(img: np.ndarray, printed_ocr: PrintedOCR) -> str:
    try:
        return printed_ocr.recognize(img).strip()
    except Exception:
        return ""


def _ollama_vision_ocr(path: Path) -> str:
    import base64
    import urllib.error
    import urllib.request
    import os

    try:
        if os.environ.get("GEMINI_API_KEY"):
            b64 = base64.b64encode(path.read_bytes()).decode("utf-8")
            gemini_resp = _gemini_generate("This is a slice of handwritten text. Please transcribe the handwriting in this image word-for-word exactly without adding any comments.", b64)
            if gemini_resp:
                return gemini_resp
        img = cv2.imread(str(path))
        if img is None:
            return ""

        h, w, c = img.shape
        if h > w and h >= 300:
            slice_h = h // 3
            texts = []
            for i in range(3):
                start = i * slice_h
                end = (i + 1) * slice_h if i < 2 else h
                sliced = img[start:end, :]

                _, buffer = cv2.imencode(".png", sliced)
                b64 = base64.b64encode(buffer).decode("utf-8")

                payload = {
                    "model": settings.vision_model,
                    "prompt": "This is a slice of handwritten text. Please transcribe the handwriting in this image word-for-word exactly without adding any comments.",
                    "images": [b64],
                    "stream": False,
                    "options": {"temperature": 0.0},
                }
                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=120) as response:
                    res = json.loads(response.read().decode("utf-8"))
                    text = res.get("response", "").strip()

                is_desc = any(phrase in text.lower() for phrase in ["image shows", "shows a page", "written in", "notebook or journal"])
                if not text or is_desc:
                    payload["prompt"] = "Transcribe the text in this image word-for-word. Output only the text."
                    data = json.dumps(payload).encode("utf-8")
                    req_fb = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"})
                    with urllib.request.urlopen(req_fb, timeout=120) as response_fb:
                        res_fb = json.loads(response_fb.read().decode("utf-8"))
                        text = res_fb.get("response", "").strip()

                if text:
                    if text.startswith('"') and text.endswith('"'):
                        text = text[1:-1].strip()
                    texts.append(text)

            return "\n\n".join(texts)
        else:
            b64 = base64.b64encode(path.read_bytes()).decode("utf-8")
            payload = {
                "model": settings.vision_model,
                "prompt": "This is a slice of handwritten text. Please transcribe the handwriting in this image word-for-word exactly without adding any comments.",
                "images": [b64],
                "stream": False,
                "options": {"temperature": 0.0},
            }
            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=120) as response:
                res = json.loads(response.read().decode("utf-8"))
                text = res.get("response", "").strip()

            is_desc = any(phrase in text.lower() for phrase in ["image shows", "shows a page", "written in", "notebook or journal"])
            if not text or is_desc:
                payload["prompt"] = "Transcribe the text in this image word-for-word. Output only the text."
                data = json.dumps(payload).encode("utf-8")
                req_fb = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"})
                with urllib.request.urlopen(req_fb, timeout=120) as response_fb:
                    res_fb = json.loads(response_fb.read().decode("utf-8"))
                    text = res_fb.get("response", "").strip()

            if text.startswith('"') and text.endswith('"'):
                text = text[1:-1].strip()
            return text
    except urllib.error.HTTPError as e:
        body = e.read().decode("utf-8", errors="replace")
        if "does not support image input" not in body:
            print(f"Ollama vision HTTP error: {body}")
    except Exception as e:
        print(f"Ollama vision failed: {e}")
    return ""


def _gemini_generate(prompt: str, base64_image: str = None, image_mime: str = "image/png") -> str:
    import os
    import json
    import urllib.request
    import urllib.error

    api_key = os.environ.get("GEMINI_API_KEY")
    if not api_key:
        return ""

    url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent?key={api_key}"
    parts = [{"text": prompt}]
    if base64_image:
        parts.append({
            "inlineData": {
                "mimeType": image_mime,
                "data": base64_image
            }
        })

    payload = {
        "contents": [{"parts": parts}]
    }

    try:
        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:
            res = json.loads(response.read().decode("utf-8"))
            return res["candidates"][0]["content"]["parts"][0]["text"].strip()
    except Exception as e:
        print(f"Gemini API call failed: {e}")
        return ""


def _llm_complete_local(prompt: str, system: str = "", temperature: float = 0.3, max_tokens: int = 512) -> str:
    import urllib.request
    import json

    gemini_resp = _gemini_generate(f"{system}\n\n{prompt}" if system else prompt)
    if gemini_resp:
        return gemini_resp

    payload = {
        "model": settings.slm_model,
        "prompt": f"{system}\n\n{prompt}" if system else prompt,
        "stream": False,
        "options": {"temperature": temperature, "num_predict": max_tokens},
    }
    url = f"{settings.ollama_url}/api/generate"
    try:
        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=90) as response:
            res = json.loads(response.read().decode("utf-8"))
            return res.get("response", "").strip()
    except Exception as e:
        print(f"Local LLM complete failed: {e}")
        return ""


def _ollama_vision_explain(path: Path) -> str:
    import base64
    import urllib.request
    import json

    try:
        b64 = base64.b64encode(path.read_bytes()).decode("utf-8")
        gemini_resp = _gemini_generate("Explain this image in very detail. Describe the layout, objects, text, colors, shapes, and overall context. Be as thorough as possible.", b64)
        if gemini_resp:
            return gemini_resp
    except Exception:
        pass

    model = settings.vision_model
    try:
        req = urllib.request.Request(f"{settings.ollama_url}/api/tags")
        with urllib.request.urlopen(req, timeout=5) as response:
            data = json.loads(response.read().decode("utf-8"))
            available_models = [m["name"] for m in data.get("models", [])]
            if model not in available_models:
                if f"{model}:latest" in available_models:
                    model = f"{model}:latest"
                elif "moondream:latest" in available_models:
                    model = "moondream:latest"
                elif "llava:7b" in available_models:
                    model = "llava:7b"
    except Exception:
        pass

    try:
        b64 = base64.b64encode(path.read_bytes()).decode("utf-8")
        payload = {
            "model": model,
            "prompt": "Explain this image in very detail. Describe the layout, objects, text, colors, shapes, and overall context. Be as thorough as possible.",
            "images": [b64],
            "stream": False,
            "options": {"temperature": 0.2},
        }
        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=120) as response:
            res = json.loads(response.read().decode("utf-8"))
            return res.get("response", "").strip()
    except Exception as e:
        print(f"Ollama vision explain failed: {e}")
        return ""


def _explain_audio_transcript(transcript: str, filename: str) -> str:
    if not transcript or transcript == "[inaudible]":
        return f"This is an audio file named '{filename}'. No clear speech could be transcribed, so a detailed audio explanation is not available."

    prompt = (
        "You are an expert audio analyst. Analyze the following transcript of an audio file and provide a very detailed explanation. "
        "Summarize the main topics discussed, analyze the speaker's likely intent or tone, point out key terms, and outline any action items.\n\n"
        f"Transcript:\n{transcript}\n\n"
        "Detailed Explanation:"
    )
    return _llm_complete_local(prompt, system="You are a detailed audio analysis assistant.", max_tokens=512)


def _explain_video_content(ocr_text: str, audio_transcript: str, filename: str) -> str:
    if (not ocr_text or ocr_text == "[UNCLEAR]") and (not audio_transcript or audio_transcript == "[inaudible]"):
        return f"This is a video file named '{filename}'. No speech or visual text could be extracted, so a detailed explanation is not available."

    prompt = (
        "You are an expert video analyst. Analyze the extracted visual text (OCR from video frames) "
        "and the speech transcript of a video file, and provide a very detailed explanation. "
        "Describe the main themes, key information shown on screen, topics discussed in the audio, "
        "and synthesize the visual and verbal content into a coherent description of what the video is about.\n\n"
        f"Extracted Visual Text:\n{ocr_text}\n\n"
        f"Speech Transcript:\n{audio_transcript}\n\n"
        "Detailed Explanation:"
    )
    return _llm_complete_local(prompt, system="You are a detailed video analysis assistant.", max_tokens=512)


def process_image(path: Path, language: str = "English") -> tuple[str, str, str, list[dict], list[list[str]]]:
    uploads_dir = _get_uploads_dir()
    uploads_dir.mkdir(exist_ok=True, parents=True)

    timestamp = int(Path(path).stat().st_mtime)
    orig_name = f"original_{timestamp}_{path.name}"
    clean_name = f"cleaned_{timestamp}_{path.name}"

    orig_dest = uploads_dir / orig_name
    clean_dest = uploads_dir / clean_name

    if path.resolve() != orig_dest.resolve():
        shutil.copy(str(path), str(orig_dest))

    img = load_image(path)
    printed_ocr = PrintedOCR()
    barcodes = []
    tables = []

    try:
        barcodes = printed_ocr.detect_barcodes(img)
    except Exception:
        pass

    try:
        tables = printed_ocr.detect_tables(img)
    except Exception:
        pass

    explanation = _ollama_vision_explain(path)

    raw_text = _ollama_vision_ocr(path)
    if raw_text:
        cleaned = preprocess(img)
        cv2.imwrite(str(clean_dest), cleaned)
    else:
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) if img.ndim == 3 else img

        strategies = [
            ("adaptive + denoise", pp_module.preprocess),
            ("adaptive (auto-tuned)", lambda x: adaptive_preprocess(x)[0]),
            ("otsu", pp_module.preprocess_otsu),
            ("clahe + otsu", pp_module.preprocess_light),
            ("grayscale", lambda x: gray.copy()),
            ("grayscale 2x", lambda x: pp_module.preprocess_with_scale(img, 2.0)),
            ("grayscale 3x", lambda x: pp_module.preprocess_with_scale(img, 3.0)),
        ]

        results = []
        for name, fn in strategies:
            try:
                cleaned = fn(img)
                text = _ocr_image(cleaned, printed_ocr)
                if text:
                    results.append((len(text), text, cleaned))
            except Exception:
                pass

        best_cleaned = None
        if results:
            results.sort(key=lambda x: x[0], reverse=True)
            _, raw_text, best_cleaned = results[0]
        else:
            raw_text = ""

        if best_cleaned is not None:
            cv2.imwrite(str(clean_dest), best_cleaned)
        else:
            cleaned_default = pp_module.preprocess(img)
            cv2.imwrite(str(clean_dest), cleaned_default)

        if not raw_text.strip():
            raw_text = _try_handwriting_ocr(best_cleaned or pp_module.preprocess(img))

        if not raw_text.strip():
            raw_text = "[UNCLEAR]"

    combined_parts = []
    if explanation:
        combined_parts.append("=========================\nIMAGE ANALYSIS & EXPLANATION\n=========================\n" + explanation)
    if raw_text and raw_text != "[UNCLEAR]":
        combined_parts.append("=========================\nEXTRACTED TEXT\n=========================\n" + raw_text)

    combined_text = "\n\n".join(combined_parts) if combined_parts else "[UNCLEAR]"
    return (combined_text, f"/static/uploads/{orig_name}", f"/static/uploads/{clean_name}", barcodes, tables)


def process_audio(path: Path, language: str = "English") -> tuple[str, str]:
    uploads_dir = _get_uploads_dir()
    uploads_dir.mkdir(exist_ok=True, parents=True)

    timestamp = int(Path(path).stat().st_mtime)
    dest_name = f"audio_{timestamp}_{path.name}"
    dest_path = uploads_dir / dest_name

    if path.resolve() != dest_path.resolve():
        shutil.copy(str(path), str(dest_path))

    transcript = ""
    try:
        transcript = _transcribe_whisper(path)
    except Exception as e:
        print(f"Whisper transcription failed, falling back to Sphinx: {e}")
        try:
            transcript = _transcribe_sphinx(path)
        except Exception as e2:
            print(f"Sphinx also failed: {e2}")

    if not transcript:
        transcript = "[inaudible]"

    explanation = _explain_audio_transcript(transcript, path.name)
    combined_parts = []
    if explanation:
        combined_parts.append("=========================\nAUDIO ANALYSIS & EXPLANATION\n=========================\n" + explanation)
    if transcript and transcript != "[inaudible]":
        combined_parts.append("=========================\nSPEECH TRANSCRIPT\n=========================\n" + transcript)

    combined_text = "\n\n".join(combined_parts) if combined_parts else "[inaudible]"
    return combined_text, f"/static/uploads/{dest_name}"


def process_video(path: Path, language: str = "English") -> tuple[str, str]:
    uploads_dir = _get_uploads_dir()
    uploads_dir.mkdir(exist_ok=True, parents=True)

    timestamp = int(Path(path).stat().st_mtime)
    dest_name = f"video_{timestamp}_{path.name}"
    dest_path = uploads_dir / dest_name

    if path.resolve() != dest_path.resolve():
        shutil.copy(str(path), str(dest_path))

    transcript = ""
    try:
        transcript = _transcribe_whisper(path)
    except Exception as e:
        print(f"Whisper transcription failed, falling back to Sphinx: {e}")
        try:
            transcript = _transcribe_sphinx(path)
        except Exception as e2:
            print(f"Sphinx also failed: {e2}")

    if not transcript:
        transcript = "[inaudible]"

    ocr_texts = []
    try:
        cap = cv2.VideoCapture(str(path))
        fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
        frame_interval = int(fps * 5)
        frame_count = 0
        success = True
        printed_ocr = PrintedOCR()
        while success:
            success, frame = cap.read()
            if not success:
                break
            if frame_count % frame_interval == 0:
                try:
                    text = printed_ocr.recognize(frame)
                    if text.strip() and text.strip() not in ocr_texts:
                        ocr_texts.append(text.strip())
                except Exception:
                    pass
            frame_count += 1
        cap.release()
    except Exception:
        pass

    ocr_text = "\n".join(ocr_texts) if ocr_texts else "[UNCLEAR]"

    explanation = _explain_video_content(ocr_text, transcript, path.name)

    combined_parts = []
    if explanation:
        combined_parts.append("=========================\nVIDEO ANALYSIS & EXPLANATION\n=========================\n" + explanation)
    if ocr_texts:
        combined_parts.append("=========================\nEXTRACTED VISUAL TEXT\n=========================\n" + ocr_text)
    if transcript and transcript != "[inaudible]":
        combined_parts.append("=========================\nSPEECH TRANSCRIPT\n=========================\n" + transcript)

    combined_text = "\n\n".join(combined_parts) if combined_parts else "[UNCLEAR]"
    return combined_text, f"/static/uploads/{dest_name}"


def process_document(path: Path, language: str = "English") -> tuple[str, str]:
    uploads_dir = _get_uploads_dir()
    uploads_dir.mkdir(exist_ok=True, parents=True)

    timestamp = int(Path(path).stat().st_mtime)
    dest_name = f"doc_{timestamp}_{path.name}"
    dest_path = uploads_dir / dest_name

    if path.resolve() != dest_path.resolve():
        shutil.copy(str(path), str(dest_path))

    ext = path.suffix.lower()
    text = ""

    if ext == ".pdf":
        try:
            import pypdfium2 as pdfium
            pdf = pdfium.PdfDocument(str(path))
            text_list = []
            for page in pdf:
                textpage = page.get_textpage()
                t = textpage.get_text_range()
                if t:
                    text_list.append(t)
            text = "\n".join(text_list).strip()
        except Exception:
            pass

        if not text.strip():
            try:
                img = load_image(path)
                cleaned = preprocess(img)
                temp_dir = Path(tempfile.gettempdir())
                temp_img_path = temp_dir / f"{path.stem}_page1.jpg"
                cv2.imwrite(str(temp_img_path), cleaned)
                printed = PrintedOCR()
                text = printed.recognize(cleaned)
                if not text.strip():
                    text = _try_handwriting_ocr(cleaned)
                if temp_img_path.exists():
                    temp_img_path.unlink()
            except Exception:
                pass

        if not text.strip():
            text = "[UNCLEAR]"

    elif ext == ".docx":
        try:
            with zipfile.ZipFile(str(path)) as docx:
                xml_content = docx.read("word/document.xml")
                root = ET.fromstring(xml_content)
                namespaces = {"w": "http://schemas.openxmlformats.org/wordprocessingml/2006/main"}
                texts = [node.text for node in root.findall(".//w:t", namespaces) if node.text]
                text = " ".join(texts)
        except Exception:
            pass
    else:
        try:
            text = path.read_text(encoding="utf-8")
        except Exception:
            try:
                text = path.read_text(encoding="latin1")
            except Exception:
                pass

    if not text.strip():
        text = "[UNCLEAR]"

    return text, f"/static/uploads/{dest_name}"


def process_spreadsheet(path: Path, language: str = "English") -> tuple[str, str]:
    import csv

    uploads_dir = _get_uploads_dir()
    uploads_dir.mkdir(exist_ok=True, parents=True)

    timestamp = int(Path(path).stat().st_mtime)
    dest_name = f"sheet_{timestamp}_{path.name}"
    dest_path = uploads_dir / dest_name

    if path.resolve() != dest_path.resolve():
        shutil.copy(str(path), str(dest_path))

    text = ""
    ext = path.suffix.lower()

    if ext == ".csv":
        try:
            with open(path, encoding="utf-8") as f:
                reader = csv.reader(f)
                rows = list(reader)
                if rows:
                    text_lines = []
                    for r in rows:
                        text_lines.append(" | ".join(r))
                    text = "CSV Spreadsheet Data:\n" + "\n".join(text_lines)
        except Exception:
            pass
    elif ext in (".xlsx", ".xls"):
        try:
            import pandas as pd
            df = pd.read_excel(str(path), engine="openpyxl" if ext == ".xlsx" else "xlrd")
            text = "Spreadsheet Data:\n" + df.to_string(index=False)
        except Exception as e:
            print(f"Spreadsheet read failed: {e}")

    if not text.strip():
        text = "[UNCLEAR]"

    return text, f"/static/uploads/{dest_name}"


def process_batch(paths: list[Path], language: str = "English", schema_type: str = "todo") -> list[dict]:
    results = []
    with ThreadPoolExecutor(max_workers=settings.max_workers) as executor:
        future_map = {}
        for path in paths:
            ext = path.suffix.lower()
            image_exts = {".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".tif", ".webp"}
            if ext in image_exts:
                future = executor.submit(process_image, path, language)
                future_map[future] = ("image", path)
            elif ext in {".mp3", ".wav", ".m4a", ".ogg", ".flac"}:
                future = executor.submit(process_audio, path, language)
                future_map[future] = ("audio", path)
            elif ext in {".mp4", ".mkv", ".mov", ".avi"}:
                future = executor.submit(process_video, path, language)
                future_map[future] = ("video", path)
            elif ext in {".pdf", ".docx", ".txt"}:
                future = executor.submit(process_document, path, language)
                future_map[future] = ("document", path)
            elif ext in {".csv", ".xlsx", ".xls"}:
                future = executor.submit(process_spreadsheet, path, language)
                future_map[future] = ("spreadsheet", path)

        for future in as_completed(future_map):
            source_type, path = future_map[future]
            try:
                if source_type == "image":
                    raw_text, image_url, cleaned_url, barcodes, tables = future.result()
                    results.append({
                        "path": str(path),
                        "source_type": source_type,
                        "raw_text": raw_text,
                        "image_url": image_url,
                        "cleaned_url": cleaned_url,
                        "barcodes": barcodes,
                        "tables": tables,
                    })
                else:
                    raw_text, file_url = future.result()
                    results.append({
                        "path": str(path),
                        "source_type": source_type,
                        "raw_text": raw_text,
                        "file_url": file_url,
                    })
            except Exception as e:
                results.append({
                    "path": str(path),
                    "source_type": source_type,
                    "error": str(e),
                })

    return results