"""
Belegscanner - KI-Dokumentenanalyse
Python Flask Web App mit CORS-Unterstützung und Poweron Design
"""

from flask import Flask, render_template, request, jsonify
from flask_cors import CORS
import requests
import base64
import json
import re
import io

# PDF Support
try:
    import fitz  # PyMuPDF
    PDF_SUPPORT = True
except ImportError:
    PDF_SUPPORT = False
    print("WARNUNG: PyMuPDF nicht installiert. PDF-Support deaktiviert.")
    print("Installieren mit: pip install pymupdf")

app = Flask(__name__)
CORS(app)  # CORS für alle Routen aktivieren


# ============================================================================
# PDF Helper Functions
# ============================================================================

def _extractImagesFromPdf(pdfBytes, maxPages=5):
    """
    Extrahiert Bilder aus einem PDF.
    Gibt eine Liste von Base64-kodierten Bildern zurück.
    """
    if not PDF_SUPPORT:
        raise Exception("PDF-Support nicht verfügbar. Bitte PyMuPDF installieren.")
    
    images = []
    
    # PDF öffnen
    doc = fitz.open(stream=pdfBytes, filetype="pdf")
    
    # Anzahl der Seiten begrenzen
    numPages = min(len(doc), maxPages)
    
    for pageNum in range(numPages):
        page = doc[pageNum]
        
        # Seite als Bild rendern (höhere Auflösung für bessere OCR)
        mat = fitz.Matrix(2.0, 2.0)  # 2x Zoom für bessere Qualität
        pix = page.get_pixmap(matrix=mat)
        
        # In PNG konvertieren
        imgBytes = pix.tobytes("png")
        imgBase64 = base64.b64encode(imgBytes).decode('utf-8')
        
        images.append({
            'page': pageNum + 1,
            'base64': imgBase64,
            'width': pix.width,
            'height': pix.height
        })
    
    doc.close()
    
    return images


def _renderPdfPageAsImage(pdfBytes, pageNum=0, zoom=2.0):
    """
    Rendert eine einzelne PDF-Seite als Bild.
    """
    if not PDF_SUPPORT:
        raise Exception("PDF-Support nicht verfügbar.")
    
    doc = fitz.open(stream=pdfBytes, filetype="pdf")
    
    if pageNum >= len(doc):
        pageNum = len(doc) - 1
    
    page = doc[pageNum]
    mat = fitz.Matrix(zoom, zoom)
    pix = page.get_pixmap(matrix=mat)
    
    imgBytes = pix.tobytes("png")
    imgBase64 = base64.b64encode(imgBytes).decode('utf-8')
    
    result = {
        'base64': imgBase64,
        'width': pix.width,
        'height': pix.height,
        'page': pageNum + 1,
        'totalPages': len(doc)
    }
    
    doc.close()
    
    return result

# ============================================================================
# Model Helper Functions
# ============================================================================

def _isVisionModel(modelName):
    """
    Prüft ob ein Modell ein Vision-Modell ist basierend auf Namenskonventionen.
    Vision-Modelle enthalten typischerweise 'vision', 'vl', 'llava', 'bakllava' im Namen.
    """
    if not modelName:
        return False
    
    modelLower = modelName.lower()
    visionIndicators = ['vision', 'vl', 'llava', 'bakllava']
    
    return any(indicator in modelLower for indicator in visionIndicators)


# ============================================================================
# Routes
# ============================================================================

@app.route('/')
def _index():
    """Hauptseite mit dem Belegscanner UI"""
    return render_template('index.html')


@app.route('/api/analyze', methods=['POST'])
def _analyzeDocument():
    """
    Analysiert ein Dokument mit Ollama Vision API oder verarbeitet Text mit Non-Vision Modellen
    Erwartet: { imageBase64 (optional bei Non-Vision), prompt, ollamaUrl, modelName }
    """
    try:
        data = request.get_json()
        
        imageBase64 = data.get('imageBase64')
        prompt = data.get('prompt')
        ollamaUrl = data.get('ollamaUrl', 'http://localhost:11434')
        modelName = data.get('modelName', 'qwen2.5vl:72b')
        
        # Prüfe ob es ein Vision-Modell ist (basierend auf Namenskonvention)
        isVisionModel = _isVisionModel(modelName)
        
        # Bei Vision-Modellen ist ein Bild erforderlich
        if isVisionModel and not imageBase64:
            return jsonify({'error': 'Kein Bild übermittelt (erforderlich für Vision-Modelle)'}), 400
        
        if not prompt:
            return jsonify({'error': 'Kein Prompt übermittelt'}), 400
        
        # Request-Body erstellen
        requestBody = {
            'model': modelName,
            'prompt': prompt,
            'stream': False
        }
        
        # Bilder nur hinzufügen wenn vorhanden (für Vision-Modelle)
        if imageBase64:
            requestBody['images'] = [imageBase64]
        
        # Ollama API aufrufen (Timeout: 60 Minuten für grosse Modelle)
        response = requests.post(
            f'{ollamaUrl}/api/generate',
            json=requestBody,
            timeout=3600  # 60 Minuten
        )
        
        if response.status_code == 404:
            return jsonify({
                'error': f'Modell "{modelName}" nicht gefunden. Bitte installieren Sie es mit: ollama pull {modelName}'
            }), 404
        
        if response.status_code != 200:
            return jsonify({
                'error': f'Ollama API Fehler: {response.status_code} - {response.text[:200]}'
            }), response.status_code
        
        responseData = response.json()
        responseText = responseData.get('response', '')
        
        # Versuche JSON aus der Antwort zu extrahieren
        extractedData = None
        jsonMatch = re.search(r'\{[\s\S]*\}', responseText)
        
        if jsonMatch:
            try:
                extractedData = json.loads(jsonMatch.group())
            except json.JSONDecodeError:
                # JSON-ähnlicher Text gefunden, aber ungültig
                extractedData = None
        
        # Wenn kein JSON gefunden, Antwort in JSON-Objekt verpacken
        if extractedData is None:
            extractedData = {
                'response': responseText.strip()
            }
        
        return jsonify({
            'success': True,
            'data': extractedData,
            'rawResponse': responseText
        })
        
    except requests.exceptions.Timeout:
        return jsonify({'error': 'Zeitüberschreitung bei der Ollama API'}), 504
    except requests.exceptions.ConnectionError:
        return jsonify({'error': 'Verbindung zu Ollama fehlgeschlagen. Ist Ollama gestartet?'}), 503
    except json.JSONDecodeError as e:
        return jsonify({'error': f'JSON Parse-Fehler: {str(e)}'}), 400
    except Exception as e:
        return jsonify({'error': f'Unerwarteter Fehler: {str(e)}'}), 500


@app.route('/api/health', methods=['GET'])
def _healthCheck():
    """Health Check Endpoint"""
    return jsonify({'status': 'ok', 'service': 'belegscanner', 'pdfSupport': PDF_SUPPORT})


@app.route('/api/pdf/extract', methods=['POST'])
def _extractPdfImages():
    """
    Extrahiert Bilder aus einem PDF.
    Erwartet: { pdfBase64, page (optional, default: alle) }
    """
    if not PDF_SUPPORT:
        return jsonify({
            'error': 'PDF-Support nicht verfügbar. Bitte PyMuPDF installieren: pip install pymupdf'
        }), 501
    
    try:
        data = request.get_json()
        pdfBase64 = data.get('pdfBase64')
        pageNum = data.get('page')  # Optional: spezifische Seite
        
        if not pdfBase64:
            return jsonify({'error': 'Kein PDF übermittelt'}), 400
        
        # Base64 dekodieren
        pdfBytes = base64.b64decode(pdfBase64)
        
        if pageNum is not None:
            # Einzelne Seite extrahieren
            result = _renderPdfPageAsImage(pdfBytes, pageNum - 1)  # 0-basiert
            return jsonify({
                'success': True,
                'image': result
            })
        else:
            # Alle Seiten extrahieren (max 5)
            images = _extractImagesFromPdf(pdfBytes, maxPages=5)
            return jsonify({
                'success': True,
                'images': images,
                'totalExtracted': len(images)
            })
        
    except Exception as e:
        return jsonify({'error': f'PDF-Verarbeitungsfehler: {str(e)}'}), 500


@app.route('/api/ollama/status', methods=['GET'])
def _ollamaStatus():
    """Prüft ob Ollama erreichbar ist und listet verfügbare Modelle"""
    ollamaUrl = request.args.get('url', 'http://localhost:11434')
    
    try:
        # Prüfe ob Ollama läuft
        response = requests.get(f'{ollamaUrl}/api/tags', timeout=5)
        
        if response.status_code != 200:
            return jsonify({
                'connected': False,
                'error': f'Ollama antwortet mit Status {response.status_code}'
            })
        
        data = response.json()
        models = [m.get('name', '') for m in data.get('models', [])]
        
        # Filtere Vision-Modelle (enthalten oft 'vision', 'vl', 'llava' im Namen)
        visionModels = [m for m in models if any(x in m.lower() for x in ['vision', 'vl', 'llava', 'bakllava'])]
        
        return jsonify({
            'connected': True,
            'models': models,
            'visionModels': visionModels,
            'totalModels': len(models)
        })
        
    except requests.exceptions.ConnectionError:
        return jsonify({
            'connected': False,
            'error': 'Keine Verbindung zu Ollama. Ist Ollama gestartet?'
        })
    except Exception as e:
        return jsonify({
            'connected': False,
            'error': str(e)
        })


# ============================================================================
# Main
# ============================================================================

if __name__ == '__main__':
    print("\n" + "="*60)
    print("  Belegscanner - KI-Dokumentenanalyse")
    print("  Powered by Poweron")
    print("="*60)
    print("\n  Server läuft auf: http://localhost:5000")
    print("  CORS ist aktiviert für alle Origins")
    print("\n  Drücke Ctrl+C zum Beenden")
    print("="*60 + "\n")
    
    app.run(host='0.0.0.0', port=5000, debug=True)