"""
Erweitertes Test-Skript für den ChatManager-Workflow mit simulierten Datei-Uploads.
Bietet zusätzliche Konfigurationsmöglichkeiten und detailliertere Tests.
"""

import asyncio
import logging
import os
import sys
import argparse
import json
from typing import Dict, Any, List, Tuple, Optional
from datetime import datetime

# Logging konfigurieren
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[logging.StreamHandler()]
)
logger = logging.getLogger("test_workflow")

# Pfad zum Projektverzeichnis hinzufügen
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# Module importieren
from modules.lucydom_interface import get_lucydom_interface
from modules.chat import get_chat_manager

class TestConfig:
    """Konfigurationsklasse für Testparameter"""
    def __init__(self):
        self.mandate_id = 1
        self.user_id = 1
        self.cleanup = True
        self.save_results = True
        self.results_dir = "test_results"
        self.test_message = "Analysiere bitte die hochgeladenen Dateien und erkläre mir deren Inhalt."
        self.text_file_content = """
        Dies ist eine Test-Textdatei für den ChatManager-Workflow.
        Sie enthält einige Informationen zum Testen der Dokumentverarbeitung.
        
        Der ChatManager sollte in der Lage sein, diese Datei zu verarbeiten
        und daraus relevante Informationen zu extrahieren.
        
        Diese Datei dient als Beispiel für Text-basierte Dokumente, die in einem
        Chat-Workflow verwendet werden können.
        """

def parse_args() -> TestConfig:
    """Parst Kommandozeilenargumente"""
    parser = argparse.ArgumentParser(description="Test für ChatManager-Workflow")
    parser.add_argument("--mandate-id", type=int, default=1, help="ID des Mandanten")
    parser.add_argument("--user-id", type=int, default=1, help="ID des Benutzers")
    parser.add_argument("--no-cleanup", action="store_true", help="Testdateien nicht löschen")
    parser.add_argument("--no-save", action="store_true", help="Ergebnisse nicht speichern")
    parser.add_argument("--results-dir", type=str, default="test_results", help="Verzeichnis für Ergebnisse")
    parser.add_argument("--message", type=str, help="Benutzernachricht für den Test")
    
    args = parser.parse_args()
    
    config = TestConfig()
    config.mandate_id = args.mandate_id
    config.user_id = args.user_id
    config.cleanup = not args.no_cleanup
    config.save_results = not args.no_save
    config.results_dir = args.results_dir
    if args.message:
        config.test_message = args.message
    
    return config

async def create_test_files(config: TestConfig) -> Tuple[int, int]:
    """
    Erstellt eine Textdatei und ein Bild für Tests und lädt sie in die Datenbank hoch.
    
    Args:
        config: Testkonfiguration
        
    Returns:
        Tuple mit (text_file_id, image_file_id)
    """
    logger.info("Erstelle Test-Dateien...")
    
    lucy_interface = get_lucydom_interface(config.mandate_id, config.user_id)
    
    # Textdatei erstellen
    text_content = config.text_file_content
    text_file_bytes = text_content.encode('utf-8')
    text_file = lucy_interface.save_uploaded_file(text_file_bytes, "test_document.txt")
    text_file_id = text_file["id"]
    logger.info(f"Textdatei erstellt mit ID: {text_file_id}")
    
    # Bilddatei erstellen (einfaches 1x1 PNG)
    # Base64-kodiertes 1x1 PNG
    png_data = bytes.fromhex(
        "89504e470d0a1a0a0000000d49484452000000010000000108060000001f15c4"
        "89000000017352474200aece1ce90000000467414d410000b18f0bfc61050000"
        "000970485973000016250000162501495224f00000001974455874536f667477"
        "617265007777772e696e6b73636170652e6f72679bee3c1a0000000c49444154"
        "08d763f8ffff3f0005fe02fec1cd59830000000049454e44ae426082"
    )
    image_file = lucy_interface.save_uploaded_file(png_data, "test_image.png")
    image_file_id = image_file["id"]
    logger.info(f"Bilddatei erstellt mit ID: {image_file_id}")
    
    return text_file_id, image_file_id

async def verify_uploaded_files(mandate_id: int, user_id: int, file_ids: List[int]) -> bool:
    """
    Überprüft, ob die hochgeladenen Dateien korrekt in der Datenbank gespeichert wurden
    
    Args:
        mandate_id: ID des Mandanten
        user_id: ID des Benutzers
        file_ids: Liste der Datei-IDs
        
    Returns:
        True, wenn alle Dateien verfügbar sind
    """
    logger.info("Überprüfe hochgeladene Dateien...")
    
    lucy_interface = get_lucydom_interface(mandate_id, user_id)
    all_files_available = True
    
    for file_id in file_ids:
        file = lucy_interface.get_file(file_id)
        if file:
            file_data = lucy_interface.get_file_data(file_id)
            if file_data:
                logger.info(f"Datei {file_id} ({file.get('name', 'Unbekannt')}, {file.get('mime_type', 'Unbekannt')}) ist verfügbar")
                logger.info(f"  Größe: {len(file_data)} Bytes")
            else:
                logger.error(f"Datei {file_id} hat keine Binärdaten")
                all_files_available = False
        else:
            logger.error(f"Datei mit ID {file_id} nicht in der Datenbank gefunden")
            all_files_available = False
    
    return all_files_available

async def run_chat_workflow(config: TestConfig, file_ids: List[int]) -> Dict[str, Any]:
    """
    Führt einen Chat-Workflow mit gegebenen Datei-IDs aus.
    
    Args:
        config: Testkonfiguration
        file_ids: Liste der Datei-IDs
        
    Returns:
        Das Workflow-Ergebnis
    """
    logger.info(f"Starte Chat-Workflow mit Dateien: {file_ids}")
    
    # ChatManager initialisieren
    chat_manager = get_chat_manager(config.mandate_id, config.user_id)
    
    # Benutzeranfrage erstellen
    user_input = {
        "message": config.test_message,
        "additional_fileids": file_ids
    }
    
    # Start-Zeit erfassen
    start_time = datetime.now()
    
    # Chat-Workflow ausführen
    workflow_result = await chat_manager.chat_run(user_input)
    
    # Ende-Zeit und Dauer berechnen
    end_time = datetime.now()
    duration = (end_time - start_time).total_seconds()
    
    logger.info(f"Workflow abgeschlossen mit ID: {workflow_result['id']}")
    logger.info(f"Dauer: {duration:.2f} Sekunden")
    
    return workflow_result

def analyze_workflow_result(workflow: Dict[str, Any]) -> Dict[str, Any]:
    """
    Analysiert das Workflow-Ergebnis und gibt Statistiken zurück.
    
    Args:
        workflow: Das Workflow-Ergebnis
        
    Returns:
        Dictionary mit Analyseergebnissen
    """
    logger.info("Analysiere Workflow-Ergebnis:")
    
    # Basis-Informationen
    analysis = {
        "workflow_id": workflow.get("id"),
        "status": workflow.get("status"),
        "message_count": len(workflow.get("messages", [])),
        "log_count": len(workflow.get("logs", [])),
        "document_count": 0,
        "roles": {},
        "document_types": {},
        "response_sizes": []
    }
    
    # Nachrichten analysieren
    for message in workflow.get("messages", []):
        # Rollen zählen
        role = message.get("role", "unknown")
        if role not in analysis["roles"]:
            analysis["roles"][role] = 0
        analysis["roles"][role] += 1
        
        # Content-Größe bei Antworten
        if role == "assistant":
            content = message.get("content", "")
            analysis["response_sizes"].append(len(content))
        
        # Dokumente zählen und analysieren
        documents = message.get("documents", [])
        analysis["document_count"] += len(documents)
        
        for doc in documents:
            contents = doc.get("contents", [])
            for content in contents:
                content_type = content.get("content_type", "unknown")
                if content_type not in analysis["document_types"]:
                    analysis["document_types"][content_type] = 0
                analysis["document_types"][content_type] += 1
    
    # Ausgabe für Log
    logger.info(f"Workflow-ID: {analysis['workflow_id']}")
    logger.info(f"Status: {analysis['status']}")
    logger.info(f"Anzahl Nachrichten: {analysis['message_count']}")
    logger.info(f"Anzahl Dokumente: {analysis['document_count']}")
    logger.info(f"Rollen-Verteilung: {analysis['roles']}")
    logger.info(f"Dokumenttypen: {analysis['document_types']}")
    
    if analysis["response_sizes"]:
        avg_size = sum(analysis["response_sizes"]) / len(analysis["response_sizes"])
        logger.info(f"Durchschnittliche Antwortgröße: {avg_size:.2f} Zeichen")
    
    # Detaillierte Nachrichteninformationen
    for i, message in enumerate(workflow.get("messages", [])[:5]):  # Begrenzung auf 5 Nachrichten
        logger.info(f"Nachricht {i+1}:")
        logger.info(f"  Rolle: {message.get('role', 'unbekannt')}")
        
        # Nur die ersten 100 Zeichen des Inhalts anzeigen
        content = message.get("content", "")
        content_preview = content[:100] + "..." if len(content) > 100 else content
        logger.info(f"  Inhalt: {content_preview}")
        
        # Dokumente in der Nachricht anzeigen
        documents = message.get("documents", [])
        if documents:
            logger.info(f"  Dokumente: {len(documents)}")
            for j, doc in enumerate(documents):
                file_id = doc.get("file_id", "keine file_id")
                logger.info(f"    Dokument {j+1}: File-ID={file_id}")
    
    return analysis

def save_test_results(config: TestConfig, workflow: Dict[str, Any], analysis: Dict[str, Any]) -> None:
    """
    Speichert die Testergebnisse in einer Datei.
    
    Args:
        config: Testkonfiguration
        workflow: Das vollständige Workflow-Ergebnis
        analysis: Die Analyseergebnisse
    """
    if not config.save_results:
        return
    
    # Ergebnisverzeichnis erstellen, falls es nicht existiert
    os.makedirs(config.results_dir, exist_ok=True)
    
    # Zeitstempel für eindeutige Dateinamen
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    
    # Speichere die Analyse
    analysis_file = os.path.join(config.results_dir, f"analysis_{timestamp}.json")
    with open(analysis_file, "w", encoding="utf-8") as f:
        json.dump(analysis, f, indent=2, ensure_ascii=False)
    logger.info(f"Analyse gespeichert in: {analysis_file}")
    
    # Speichere den vollständigen Workflow (ohne große Binärdaten)
    workflow_copy = workflow.copy()
    
    # Entferne Binärdaten aus dem Export, um die Dateigröße zu reduzieren
    for message in workflow_copy.get("messages", []):
        if "documents" in message:
            for doc in message.get("documents", []):
                if "contents" in doc:
                    for content in doc.get("contents", []):
                        if "data" in content and isinstance(content["data"], bytes) and len(content["data"]) > 1000:
                            content["data"] = f"[{len(content['data'])} Bytes]"
    
    workflow_file = os.path.join(config.results_dir, f"workflow_{timestamp}.json")
    with open(workflow_file, "w", encoding="utf-8") as f:
        # Konvertiere Bytes zu Strings für JSON-Serialisierung
        json.dump(workflow_copy, f, indent=2, ensure_ascii=False, default=lambda o: 
                  o.decode("utf-8") if isinstance(o, bytes) else str(o))
    logger.info(f"Workflow gespeichert in: {workflow_file}")

async def cleanup_test_files(config: TestConfig, file_ids: List[int]) -> None:
    """
    Bereinigt die erstellten Testdateien.
    
    Args:
        config: Testkonfiguration
        file_ids: Liste der zu löschenden Datei-IDs
    """
    if not config.cleanup:
        logger.info("Bereinigung übersprungen (--no-cleanup)")
        return
    
    logger.info("Beginne Bereinigung der Testdateien...")
    
    lucy_interface = get_lucydom_interface(config.mandate_id, config.user_id)
    
    for file_id in file_ids:
        try:
            success = lucy_interface.delete_file(file_id)
            if success:
                logger.info(f"Datei mit ID {file_id} erfolgreich gelöscht")
            else:
                logger.warning(f"Fehler beim Löschen der Datei mit ID {file_id}")
        except Exception as e:
            logger.error(f"Fehler beim Löschen der Datei mit ID {file_id}: {str(e)}")
    
    logger.info("Bereinigung abgeschlossen")

async def main():
    """
    Hauptfunktion, die den gesamten Testprozess steuert.
    """
    # Konfiguration laden
    config = parse_args()
    
    try:
        logger.info("=== Test-Workflow für ChatManager gestartet ===")
        logger.info(f"Mandate-ID: {config.mandate_id}, User-ID: {config.user_id}")
        
        # Schritt 1: Testdateien erstellen
        text_file_id, image_file_id = await create_test_files(config)
        file_ids = [text_file_id, image_file_id]
        
        # Schritt 2: Hochgeladene Dateien überprüfen
        files_ok = await verify_uploaded_files(config.mandate_id, config.user_id, file_ids)
        if not files_ok:
            logger.error("Fehler bei den hochgeladenen Dateien, Test wird abgebrochen")
            return
        
        # Schritt 3: Chat-Workflow ausführen
        workflow_result = await run_chat_workflow(config, file_ids)
        
        # Schritt 4: Ergebnis analysieren
        analysis = analyze_workflow_result(workflow_result)
        
        # Schritt 5: Ergebnisse speichern
        save_test_results(config, workflow_result, analysis)
        
        # Schritt 6: Bereinigen
        await cleanup_test_files(config, file_ids)
        
        logger.info("=== Test-Workflow erfolgreich abgeschlossen ===")
        
    except Exception as e:
        logger.error(f"Fehler im Test-Workflow: {str(e)}", exc_info=True)
        logger.info("=== Test-Workflow mit Fehler beendet ===")

if __name__ == "__main__":
    # Event-Loop für asyncio erstellen und Hauptfunktion ausführen
    loop = asyncio.get_event_loop()
    loop.run_until_complete(main())