#!/usr/bin/env python3
"""
Testskript zum Erstellen eines Workflows mit Prompt und Datei über den Gateway.
"""

import requests
import json
import os
import time
import sys
from datetime import datetime

# Konfiguration
API_BASE_URL = "http://localhost:8000"  # Anpassen an deine Gateway-URL
API_TOKEN = "your_api_token_here"  # Dein API-Token

# Headers für Authentifizierung
HEADERS = {
    "Authorization": f"Bearer {API_TOKEN}",
    "Content-Type": "application/json"
}

def log_message(message):
    """Gibt eine formatierte Nachricht mit Zeitstempel aus"""
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    print(f"[{timestamp}] {message}")

def upload_file(file_path):
    """Lädt eine Datei hoch und gibt die Datei-ID zurück"""
    log_message(f"Lade Datei hoch: {file_path}")
    
    if not os.path.exists(file_path):
        log_message(f"FEHLER: Datei nicht gefunden: {file_path}")
        return None
    
    # Multipart-Formular für Datei-Upload vorbereiten
    file_name = os.path.basename(file_path)
    files = {
        'file': (file_name, open(file_path, 'rb'), 'application/octet-stream')
    }
    
    # Datei hochladen
    upload_url = f"{API_BASE_URL}/api/files/upload"
    response = requests.post(
        upload_url,
        headers={"Authorization": f"Bearer {API_TOKEN}"},  # Nur Authorization-Header
        files=files
    )
    
    if response.status_code != 200:
        log_message(f"FEHLER: Datei-Upload fehlgeschlagen. Status: {response.status_code}")
        log_message(f"Response: {response.text}")
        return None
    
    # Datei-ID extrahieren
    file_data = response.json()
    file_id = file_data.get("id")
    log_message(f"Datei erfolgreich hochgeladen. ID: {file_id}")
    
    return file_id

def create_workflow(prompt, file_id=None):
    """Erstellt einen neuen Workflow mit Prompt und optionaler Datei"""
    log_message("Erstelle neuen Workflow...")
    
    # Nachricht für den Workflow vorbereiten
    user_input = {
        "message": prompt
    }
    
    # Wenn eine Datei-ID vorhanden ist, füge sie hinzu
    if file_id:
        user_input["additional_files"] = [file_id]
    
    # Workflow erstellen
    workflow_url = f"{API_BASE_URL}/api/workflows/user-input"
    response = requests.post(
        workflow_url,
        headers=HEADERS,
        json=user_input
    )
    
    if response.status_code >= 400:
        log_message(f"FEHLER: Workflow-Erstellung fehlgeschlagen. Status: {response.status_code}")
        log_message(f"Response: {response.text}")
        return None
    
    # Workflow-ID extrahieren
    workflow_data = response.json()
    workflow_id = workflow_data.get("workflow_id")
    log_message(f"Workflow erfolgreich erstellt. ID: {workflow_id}")
    
    return workflow_id

def poll_workflow_status(workflow_id, max_attempts=20, delay=2):
    """Fragt den Status eines Workflows ab und wartet bis zur Fertigstellung"""
    log_message(f"Prüfe Status des Workflows {workflow_id}...")
    
    for attempt in range(1, max_attempts + 1):
        status_url = f"{API_BASE_URL}/api/workflows/{workflow_id}/status"
        response = requests.get(
            status_url,
            headers=HEADERS
        )
        
        if response.status_code != 200:
            log_message(f"FEHLER: Status-Abfrage fehlgeschlagen. Status: {response.status_code}")
            continue
        
        status_data = response.json()
        current_status = status_data.get("status")
        log_message(f"Workflow-Status: {current_status} (Versuch {attempt}/{max_attempts})")
        
        if current_status in ["completed", "stopped", "failed"]:
            log_message(f"Workflow ist abgeschlossen mit Status: {current_status}")
            return status_data
        
        time.sleep(delay)
    
    log_message(f"Maximale Anzahl von Versuchen erreicht. Letzter Status: {current_status}")
    return None

def get_workflow_messages(workflow_id):
    """Ruft alle Nachrichten eines Workflows ab"""
    log_message(f"Hole Nachrichten für Workflow {workflow_id}...")
    
    messages_url = f"{API_BASE_URL}/api/workflows/{workflow_id}/messages"
    response = requests.get(
        messages_url,
        headers=HEADERS
    )
    
    if response.status_code != 200:
        log_message(f"FEHLER: Abrufen der Nachrichten fehlgeschlagen. Status: {response.status_code}")
        return []
    
    messages = response.json()
    log_message(f"{len(messages)} Nachrichten gefunden.")
    
    return messages

def print_workflow_results(workflow_id):
    """Gibt die Ergebnisse eines Workflows aus"""
    log_message("=== WORKFLOW-ERGEBNISSE ===")
    
    # Status abrufen
    status_url = f"{API_BASE_URL}/api/workflows/{workflow_id}/status"
    status_response = requests.get(status_url, headers=HEADERS)
    
    if status_response.status_code == 200:
        status_data = status_response.json()
        log_message(f"Workflow-Name: {status_data.get('name')}")
        log_message(f"Status: {status_data.get('status')}")
        log_message(f"Gestartet: {status_data.get('started_at')}")
        log_message(f"Letzte Aktivität: {status_data.get('last_activity')}")
    
    # Nachrichten abrufen und ausgeben
    messages = get_workflow_messages(workflow_id)
    log_message(f"Anzahl der Nachrichten: {len(messages)}")
    
    for i, msg in enumerate(messages, 1):
        log_message(f"--- Nachricht {i} ---")
        log_message(f"Rolle: {msg.get('role')}")
        
        # Inhalt gekürzt ausgeben (maximal ersten 200 Zeichen)
        content = msg.get('content', '')
        if content:
            if len(content) > 200:
                log_message(f"Inhalt: {content[:200]}... [gekürzt]")
            else:
                log_message(f"Inhalt: {content}")
        
        # Anzahl der Dokumente ausgeben
        docs = msg.get('documents', [])
        if docs:
            log_message(f"Dokumente: {len(docs)}")
            for j, doc in enumerate(docs, 1):
                source = doc.get('source', {})
                doc_name = source.get('name', f"Dokument {j}")
                log_message(f"  - {doc_name}")

def main():
    """Hauptfunktion zum Testen des Workflows"""
    # Beispiel-Datei zum Hochladen (Pfad anpassen)
    file_path = "example.csv"  # Hier den Pfad zu deiner Testdatei angeben
    
    # Prompt für den Workflow
    test_prompt = """Bitte analysiere die angehängte Datei und erstelle eine Zusammenfassung der wichtigsten Informationen.
    Wenn es sich um eine CSV-Datei handelt, identifiziere die Spalten und gib mir einen Überblick über die enthaltenen Daten.
    Erstelle außerdem eine Visualisierung, wenn du Zahlenwerte in der Datei findest."""
    
    try:
        # Datei hochladen
        file_id = upload_file(file_path)
        if not file_id:
            log_message("Test abgebrochen: Datei konnte nicht hochgeladen werden.")
            return False
        
        # Workflow erstellen
        workflow_id = create_workflow(test_prompt, file_id)
        if not workflow_id:
            log_message("Test abgebrochen: Workflow konnte nicht erstellt werden.")
            return False
        
        # Auf Abschluss des Workflows warten
        workflow_status = poll_workflow_status(workflow_id)
        if not workflow_status:
            log_message("Test unvollständig: Timeout beim Warten auf Workflow-Abschluss.")
            # Weiter machen, um zumindest Teilergebnisse zu sehen
        
        # Ergebnisse ausgeben
        print_workflow_results(workflow_id)
        
        log_message("Test abgeschlossen.")
        return True
        
    except Exception as e:
        log_message(f"FEHLER: Unerwartete Ausnahme: {str(e)}")
        import traceback
        log_message(traceback.format_exc())
        return False

if __name__ == "__main__":
    log_message("=== WORKFLOW-TEST GESTARTET ===")
    success = main()
    log_message(f"=== WORKFLOW-TEST BEENDET (Erfolgreich: {success}) ===")
    sys.exit(0 if success else 1)