#!/usr/bin/env python3
# Copyright (c) 2025 Patrick Motsch
# All rights reserved.
"""
Integration tests for workflow execution
Tests full workflow execution with state management, Stage 1/2, document extraction flow.
"""

import pytest
import uuid
from unittest.mock import Mock, AsyncMock, patch

from modules.aichat.datamodelFeatureAiChat import ChatWorkflow, TaskContext, TaskStep
from modules.datamodels.datamodelWorkflow import ActionDefinition
from modules.datamodels.datamodelDocref import DocumentReferenceList, DocumentListReference, DocumentItemReference


class TestWorkflowStateManagement:
    """Test workflow state management during execution"""
    
    @pytest.mark.asyncio
    async def test_workflow_state_increments(self):
        """Test that workflow state increments correctly during execution"""
        workflow = ChatWorkflow(
            id=str(uuid.uuid4()),
            name="Test Workflow",
            mandateId="test_mandate"
        )
        
        # Initial state
        assert workflow.currentRound == 0
        assert workflow.currentTask == 0
        assert workflow.currentAction == 0
        
        # Simulate workflow progression
        workflow.incrementAction()
        assert workflow.currentAction == 1
        
        workflow.incrementTask()
        assert workflow.currentTask == 1
        assert workflow.currentAction == 0  # Reset when task increments
        
        workflow.incrementRound()
        assert workflow.currentRound == 1
        assert workflow.currentTask == 0  # Reset when round increments
        assert workflow.currentAction == 0


class TestStage1ToStage2Flow:
    """Test Stage 1 → Stage 2 parameter generation flow"""
    
    def test_actionDefinition_needsStage2_logic(self):
        """Test needsStage2() deterministic logic"""
        # Stage 1: No parameters
        actionDef = ActionDefinition(
            action="ai.process",
            actionObjective="Process documents"
        )
        assert actionDef.needsStage2() is True
        
        # Stage 2: Parameters added
        actionDef.parameters = {"resultType": "pdf"}
        assert actionDef.needsStage2() is False
    
    def test_actionDefinition_stage1_resources(self):
        """Test that Stage 1 always defines documentList and connectionReference if needed"""
        docList = DocumentReferenceList(references=[
            DocumentListReference(label="task1_results")
        ])
        actionDef = ActionDefinition(
            action="ai.process",
            actionObjective="Process documents",
            documentList=docList,
            connectionReference="conn123"
        )
        # Stage 1 resources are set, but parameters are not
        assert actionDef.documentList is not None
        assert actionDef.connectionReference == "conn123"
        assert actionDef.needsStage2() is True  # Still needs Stage 2 for parameters


class TestDocumentExtractionFlow:
    """Test document extraction → AI processing flow"""
    
    def test_extractContentParameters_structure(self):
        """Test ExtractContentParameters structure"""
        from modules.datamodels.datamodelWorkflow import ExtractContentParameters
        
        docList = DocumentReferenceList(references=[
            DocumentListReference(label="input_docs")
        ])
        params = ExtractContentParameters(documentList=docList)
        
        assert params.documentList is not None
        assert len(params.documentList.references) == 1
        assert params.extractionOptions is None  # Optional
    
    def test_documentReferenceList_parsing(self):
        """Test DocumentReferenceList parsing from strings"""
        stringList = [
            "docList:task1_results",
            "docItem:doc123:test.pdf"
        ]
        refList = DocumentReferenceList.from_string_list(stringList)
        
        assert len(refList.references) == 2
        assert isinstance(refList.references[0], DocumentListReference)
        assert isinstance(refList.references[1], DocumentItemReference)


class TestDocumentReferenceLookup:
    """Test document reference lookup across tasks/rounds"""
    
    def test_documentListReference_with_messageId(self):
        """Test DocumentListReference with messageId for cross-round references"""
        ref = DocumentListReference(
            messageId="msg123",
            label="task1_results"
        )
        assert ref.messageId == "msg123"
        assert ref.label == "task1_results"
        assert ref.to_string() == "docList:msg123:task1_results"
    
    def test_documentListReference_without_messageId(self):
        """Test DocumentListReference without messageId (current message)"""
        ref = DocumentListReference(label="task1_results")
        assert ref.messageId is None
        assert ref.to_string() == "docList:task1_results"


class TestJsonParsing:
    """Test JSON parsing with broken/incomplete JSON"""
    
    def test_parseJsonWithModel_with_code_fences(self):
        """Test parseJsonWithModel handles code fences"""
        from modules.shared.jsonUtils import parseJsonWithModel
        
        jsonStr = '```json\n{"action": "ai.process", "actionObjective": "Process"}\n```'
        result = parseJsonWithModel(jsonStr, ActionDefinition)
        
        assert isinstance(result, ActionDefinition)
        assert result.action == "ai.process"
    
    def test_parseJsonWithModel_with_extra_text(self):
        """Test parseJsonWithModel extracts JSON from text with extra content"""
        from modules.shared.jsonUtils import parseJsonWithModel
        
        jsonStr = 'Some text before {"action": "ai.process", "actionObjective": "Process"} some text after'
        result = parseJsonWithModel(jsonStr, ActionDefinition)
        
        assert isinstance(result, ActionDefinition)
        assert result.action == "ai.process"


if __name__ == "__main__":
    pytest.main([__file__, "-v"])