engines/openwebui_engine.py

"""
OpenWebUI Engine Plugin

Wraps the OpenWebUI API (OpenAI-compatible) from uni-freiburg.de as an HTR engine.
Supports multiple models available on the OpenWebUI platform.
"""

from typing import Dict, Any, Optional, List
from pathlib import Path
import numpy as np
from PIL import Image
import io
import base64

from htr_engine_base import HTREngine, TranscriptionResult

try:
    from PyQt6.QtWidgets import (
        QWidget, QVBoxLayout, QHBoxLayout, QLabel, QComboBox,
        QPushButton, QCheckBox, QLineEdit, QGroupBox, QTextEdit,
        QSpinBox
    )
    from PyQt6.QtCore import Qt
    PYQT_AVAILABLE = True
except ImportError:
    PYQT_AVAILABLE = False
    QWidget = object

try:
    from openai import OpenAI
    OPENAI_AVAILABLE = True
except ImportError:
    OPENAI_AVAILABLE = False

try:
    from dotenv import load_dotenv
    DOTENV_AVAILABLE = True
except ImportError:
    DOTENV_AVAILABLE = False


class OpenWebUIEngine(HTREngine):
    """OpenWebUI API HTR engine plugin (OpenAI-compatible)."""

    def __init__(self):
        self.client: Optional[OpenAI] = None
        self._config_widget: Optional[QWidget] = None
        self._available_models: List[str] = []

        # Store config from load_model for batch processing
        self._loaded_config: Dict[str, Any] = {}

        # Widget references
        self._model_combo: Optional[QComboBox] = None
        self._api_key_edit: Optional[QLineEdit] = None
        self._show_key_check: Optional[QCheckBox] = None
        self._prompt_edit: Optional[QTextEdit] = None
        self._temperature_spin: Optional[QSpinBox] = None
        self._max_tokens_spin: Optional[QSpinBox] = None
        self._refresh_models_btn: Optional[QPushButton] = None

        # Default API configuration
        self.base_url = ""
        
        # Load environment variables from .env file (only once when instantiated)
        self._load_env_variables()

    def _load_env_variables(self):
        """Load environment variables from .env file if available."""
        try:
            from dotenv import load_dotenv
            # Look for .env in the project root (parent of engines/)
            env_path = Path(__file__).parent.parent / ".env"
            if env_path.exists():
                load_dotenv(env_path)
        except ImportError:
            # Silently skip if python-dotenv is not installed
            # Environment variables can still be set via OS
            pass

        # Load environment variables from .env file (if available)
        self._load_env_file()

    def _load_env_file(self):
        """Load environment variables from project root's .env file.
        
        Looks for .env in the project root directory (parent of engines/).
        Silently skips loading if python-dotenv is not installed or if .env doesn't exist.
        
        If .env loading fails or is skipped, the engine will still work if the API key
        is provided through the config dict.
        """
        if not DOTENV_AVAILABLE:
            return
            
        env_path = Path(__file__).parent.parent / ".env"
        if env_path.exists():
            load_dotenv(env_path)

    def get_name(self) -> str:
        return "OpenWebUI"

    def get_description(self) -> str:
        return "OpenWebUI API from openwebui.uni-freiburg.de (OpenAI-compatible, multiple models)"

    def is_available(self) -> bool:
        return OPENAI_AVAILABLE

    def get_unavailable_reason(self) -> str:
        if not OPENAI_AVAILABLE:
            return "OpenAI library not installed. Install with: pip install openai"
        return ""

    def get_config_widget(self) -> QWidget:
        """Create OpenWebUI configuration panel."""
        if self._config_widget is not None:
            return self._config_widget

        widget = QWidget()
        layout = QVBoxLayout()

        # API Key section
        key_group = QGroupBox("API Key")
        key_layout = QVBoxLayout()

        key_input_layout = QHBoxLayout()
        self._api_key_edit = QLineEdit()
        self._api_key_edit.setEchoMode(QLineEdit.EchoMode.Password)
        self._api_key_edit.setPlaceholderText("Enter your OpenWebUI API key")
        key_input_layout.addWidget(self._api_key_edit)

        self._show_key_check = QCheckBox("Show")
        self._show_key_check.toggled.connect(self._toggle_key_visibility)
        key_input_layout.addWidget(self._show_key_check)
        key_layout.addLayout(key_input_layout)

        key_hint = QLabel("Get your API key from https://openwebui.uni-freiburg.de")
        key_hint.setStyleSheet("color: gray; font-size: 9pt;")
        key_layout.addWidget(key_hint)

        key_group.setLayout(key_layout)
        layout.addWidget(key_group)

        # Model selection with refresh button
        model_group = QGroupBox("Model Selection")
        model_layout = QVBoxLayout()

        model_select_layout = QHBoxLayout()
        self._model_combo = QComboBox()
        self._model_combo.setMinimumWidth(300)
        model_select_layout.addWidget(self._model_combo)

        self._refresh_models_btn = QPushButton("Refresh Models")
        self._refresh_models_btn.clicked.connect(self._refresh_models)
        model_select_layout.addWidget(self._refresh_models_btn)

        model_layout.addLayout(model_select_layout)

        model_hint = QLabel("Click 'Refresh Models' to load available models from the server")
        model_hint.setStyleSheet("color: gray; font-size: 9pt;")
        model_layout.addWidget(model_hint)

        model_group.setLayout(model_layout)
        layout.addWidget(model_group)

        # Generation parameters
        params_group = QGroupBox("Generation Parameters")
        params_layout = QVBoxLayout()

        # Temperature
        temp_layout = QHBoxLayout()
        temp_layout.addWidget(QLabel("Temperature:"))
        self._temperature_spin = QSpinBox()
        self._temperature_spin.setRange(0, 100)
        self._temperature_spin.setValue(10)  # 0.1
        self._temperature_spin.setSuffix(" (×0.01)")
        temp_layout.addWidget(self._temperature_spin)
        temp_layout.addStretch()
        params_layout.addLayout(temp_layout)

        # Max tokens
        tokens_layout = QHBoxLayout()
        tokens_layout.addWidget(QLabel("Max Tokens:"))
        self._max_tokens_spin = QSpinBox()
        self._max_tokens_spin.setRange(100, 4096)
        self._max_tokens_spin.setValue(500)
        tokens_layout.addWidget(self._max_tokens_spin)
        tokens_layout.addStretch()
        params_layout.addLayout(tokens_layout)

        params_group.setLayout(params_layout)
        layout.addWidget(params_group)

        # Custom prompt section
        prompt_group = QGroupBox("Custom Prompt (Optional)")
        prompt_layout = QVBoxLayout()

        self._prompt_edit = QTextEdit()
        self._prompt_edit.setPlaceholderText(
            "Enter custom transcription prompt...\n\n"
            "Default prompt:\n"
            "Transcribe the text in this historical manuscript line image. "
            "Return only the transcribed text without any explanation or formatting."
        )
        self._prompt_edit.setMaximumHeight(120)
        prompt_layout.addWidget(self._prompt_edit)

        prompt_group.setLayout(prompt_layout)
        layout.addWidget(prompt_group)

        layout.addStretch()
        widget.setLayout(layout)

        self._config_widget = widget

        # Try to load saved API key
        self._load_saved_api_key()

        return widget

    def _toggle_key_visibility(self, checked: bool):
        """Toggle API key visibility."""
        if checked:
            self._api_key_edit.setEchoMode(QLineEdit.EchoMode.Normal)
        else:
            self._api_key_edit.setEchoMode(QLineEdit.EchoMode.Password)

    def _get_api_key_file(self) -> 'Path':
        """Get path to API key storage file."""
        from pathlib import Path
        storage_dir = Path.home() / ".trocr_gui"
        storage_dir.mkdir(exist_ok=True)
        return storage_dir / "api_keys.json"

    def _load_saved_api_key(self):
        """Load saved API key."""
        try:
            import json
            key_file = self._get_api_key_file()

            if key_file.exists():
                with open(key_file, "r") as f:
                    keys = json.load(f)

                if "openwebui" in keys:
                    self._api_key_edit.setText(keys["openwebui"])
        except Exception as e:
            print(f"Warning: Could not load saved API key: {e}")

    def _save_api_key(self):
        """Save API key."""
        try:
            import json
            key_file = self._get_api_key_file()

            # Load existing keys
            keys = {}
            if key_file.exists():
                with open(key_file, "r") as f:
                    keys = json.load(f)

            # Update key for OpenWebUI
            api_key = self._api_key_edit.text().strip()

            if api_key:
                keys["openwebui"] = api_key

                with open(key_file, "w") as f:
                    json.dump(keys, f, indent=2)
        except Exception as e:
            print(f"Warning: Could not save API key: {e}")

    def _refresh_models(self):
        """Fetch available models from OpenWebUI API."""
        api_key = self._api_key_edit.text().strip()

        if not api_key:
            self._model_combo.clear()
            self._model_combo.addItem("Please enter API key first")
            return

        try:
            # Create temporary client to fetch models
            client = OpenAI(
                base_url=self.base_url,
                api_key=api_key
            )

            # Fetch models
            models = client.models.list()

            self._available_models = []
            for model in models.data:
                self._available_models.append(model.id)

            # Update combo box
            self._model_combo.clear()
            if self._available_models:
                self._model_combo.addItems(sorted(self._available_models))
                print(f"[OpenWebUI] Loaded {len(self._available_models)} models")
            else:
                self._model_combo.addItem("No models found")

        except Exception as e:
            print(f"Error fetching models: {e}")
            self._model_combo.clear()
            self._model_combo.addItem(f"Error: {str(e)[:50]}")

    def get_config(self) -> Dict[str, Any]:
        """Extract configuration from widget controls."""
        if self._config_widget is None:
            return {}

        prompt_text = self._prompt_edit.toPlainText().strip()

        return {
            "api_key": self._api_key_edit.text().strip(),
            "model": self._model_combo.currentText(),
            "temperature": self._temperature_spin.value() / 100.0,
            "max_tokens": self._max_tokens_spin.value(),
            "custom_prompt": prompt_text if prompt_text else None,
        }

    def set_config(self, config: Dict[str, Any]):
        """Restore configuration to widget controls."""
        if self._config_widget is None:
            return

        self._api_key_edit.setText(config.get("api_key", ""))

        model = config.get("model", "")
        idx = self._model_combo.findText(model)
        if idx >= 0:
            self._model_combo.setCurrentIndex(idx)

        temp = int(config.get("temperature", 0.1) * 100)
        self._temperature_spin.setValue(temp)

        self._max_tokens_spin.setValue(config.get("max_tokens", 500))

        custom_prompt = config.get("custom_prompt", "")
        if custom_prompt:
            self._prompt_edit.setPlainText(custom_prompt)

    def load_model(self, config: Dict[str, Any]) -> bool:
        """Initialize OpenWebUI client."""
        try:
            api_key = config.get("api_key", "")

            if not api_key:
                print("Error: No API key provided. Paste your key in the field.")
                return False

            base_url = config.get("base_url", "").strip().rstrip("/")
            if not base_url:
                print("Error: No OpenWebUI base URL provided.")
                return False

            # Store config for batch processing (model, temperature, etc.)
            self._loaded_config = config.copy()

            # Save API key for future use
            if self._api_key_edit and self._api_key_edit.text().strip():
                self._save_api_key()

            self.base_url = base_url

            # Initialize client
            self.client = OpenAI(
                base_url=self.base_url,
                api_key=api_key
            )

            model = config.get("model", config.get("model_id", "unknown"))
            print(f"[OpenWebUI] Client initialized with base URL: {self.base_url}, model: {model}")
            return True

        except Exception as e:
            print(f"Error initializing OpenWebUI client: {e}")
            self.client = None
            return False

    def unload_model(self):
        """Unload OpenWebUI client."""
        if self.client is not None:
            self.client = None
        self._loaded_config = {}

    def is_model_loaded(self) -> bool:
        """Check if client is initialized."""
        return self.client is not None

    def transcribe_line(self, image: np.ndarray, config: Optional[Dict[str, Any]] = None) -> TranscriptionResult:
        """Transcribe a line image with OpenWebUI API."""
        if self.client is None:
            return TranscriptionResult(text="[OpenWebUI client not initialized]", confidence=0.0)

        if config is None:
            # First try loaded config (from batch processing), then GUI config
            if self._loaded_config:
                config = self._loaded_config
            else:
                config = self.get_config()

        try:
            # Convert numpy array to PIL Image
            if isinstance(image, np.ndarray):
                pil_image = Image.fromarray(image)
            else:
                pil_image = image

            # Convert to RGB if needed
            if pil_image.mode != 'RGB':
                pil_image = pil_image.convert('RGB')

            # Encode image to base64
            buffered = io.BytesIO()
            pil_image.save(buffered, format="PNG")
            img_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')

            # Prepare prompt
            custom_prompt = config.get("custom_prompt")
            if custom_prompt:
                prompt = custom_prompt
            else:
                prompt = (
                    "Transcribe the text in this historical manuscript line image. "
                    "Return only the transcribed text without any explanation or formatting."
                )

            # Get model and parameters
            model = config.get("model", "gpt-4-vision-preview")
            temperature = config.get("temperature", 0.1)
            max_tokens = config.get("max_tokens")
            # Treat 0 as "no limit" (HTML number fields send 0 for blank)
            if max_tokens is not None and max_tokens <= 0:
                max_tokens = None

            # Call OpenWebUI API (OpenAI-compatible)
            api_kwargs = dict(
                model=model,
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "text",
                                "text": prompt
                            },
                            {
                                "type": "image_url",
                                "image_url": {
                                    "url": f"data:image/png;base64,{img_base64}"
                                }
                            }
                        ]
                    }
                ],
                temperature=temperature,
            )
            if max_tokens is not None:
                api_kwargs["max_tokens"] = max_tokens
            response = self.client.chat.completions.create(**api_kwargs)

            # Extract transcription
            text = response.choices[0].message.content.strip()

            # Extract usage info
            usage = {}
            if hasattr(response, 'usage') and response.usage:
                usage = {
                    "prompt_tokens": response.usage.prompt_tokens,
                    "completion_tokens": response.usage.completion_tokens,
                    "total_tokens": response.usage.total_tokens
                }

            return TranscriptionResult(
                text=text,
                confidence=1.0,  # OpenWebUI doesn't provide confidence
                metadata={
                    "provider": "openwebui",
                    "model": model,
                    "usage": usage
                }
            )

        except Exception as e:
            print(f"Error in OpenWebUI transcription: {e}")
            import traceback
            traceback.print_exc()
            return TranscriptionResult(text=f"[OpenWebUI Error: {e}]", confidence=0.0)

    def get_capabilities(self) -> Dict[str, bool]:
        """OpenWebUI capabilities."""
        return {
            "batch_processing": False,
            "confidence_scores": False,
            "beam_search": False,
            "language_model": True,
            "preprocessing": True,
        }

    def requires_line_segmentation(self) -> bool:
        """OpenWebUI VLMs can process full pages directly without segmentation."""
        return False  # VLMs process full page images