# Copyright (c) 2025 Patrick Motsch
# All rights reserved.
"""
Generic Looping Use Case System

Provides parametrized looping infrastructure supporting different JSON formats and use cases.
"""

import logging
from dataclasses import dataclass, field
from typing import Dict, Any, List, Optional, Callable

logger = logging.getLogger(__name__)


@dataclass
class LoopingUseCase:
    """Configuration for a specific looping use case."""
    
    # Identification
    useCaseId: str  # "section_content", "chapter_structure", "document_structure", "code_structure", "code_content", "image_batch"
    
    # JSON Format Detection
    jsonTemplate: Dict[str, Any]  # Expected JSON structure template
    detectionKeys: List[str]  # Keys to check for format detection (e.g., ["elements"], ["chapters"], ["files"])
    detectionPath: str  # JSONPath to check (e.g., "documents[0].chapters", "files[0].content")
    
    # Prompt Building
    initialPromptBuilder: Optional[Callable] = None  # Function to build initial prompt
    continuationPromptBuilder: Optional[Callable] = None  # Function to build continuation prompt
    
    # Accumulation & Merging
    accumulator: Optional[Callable] = None  # Function to accumulate fragments
    merger: Optional[Callable] = None  # Function to merge accumulated data
    
    # Continuation Context
    continuationContextBuilder: Optional[Callable] = None  # Build continuation context for this format
    
    # Result Building
    resultBuilder: Optional[Callable] = None  # Build final result from accumulated data
    
    # Metadata
    supportsAccumulation: bool = True  # Whether this use case supports accumulation
    requiresExtraction: bool = False  # Whether this requires extraction (like sections)


class LoopingUseCaseRegistry:
    """Registry of all looping use cases."""
    
    def __init__(self):
        self.useCases: Dict[str, LoopingUseCase] = {}
        self._registerDefaultUseCases()
    
    def register(self, useCase: LoopingUseCase):
        """Register a new use case."""
        self.useCases[useCase.useCaseId] = useCase
        logger.debug(f"Registered looping use case: {useCase.useCaseId}")
    
    def get(self, useCaseId: str) -> Optional[LoopingUseCase]:
        """Get use case by ID."""
        return self.useCases.get(useCaseId)
    
    def detectUseCase(self, parsedJson: Dict[str, Any]) -> Optional[str]:
        """Detect which use case matches the JSON structure."""
        for useCaseId, useCase in self.useCases.items():
            if self._matchesFormat(parsedJson, useCase):
                return useCaseId
        return None
    
    def _matchesFormat(self, json: Dict[str, Any], useCase: LoopingUseCase) -> bool:
        """Check if JSON matches use case format."""
        # Check top-level keys
        for key in useCase.detectionKeys:
            if key in json:
                return True
        
        # Check nested path using simple dictionary traversal (no jsonpath_ng needed)
        if useCase.detectionPath:
            try:
                # Simple path matching without jsonpath_ng
                # Format: "documents[0].chapters" or "files[0].content"
                pathParts = useCase.detectionPath.split(".")
                current = json
                
                for part in pathParts:
                    # Handle array indices like "documents[0]"
                    if "[" in part and "]" in part:
                        key = part.split("[")[0]
                        index = int(part.split("[")[1].split("]")[0])
                        if isinstance(current, dict) and key in current:
                            if isinstance(current[key], list) and 0 <= index < len(current[key]):
                                current = current[key][index]
                            else:
                                return False
                        else:
                            return False
                    else:
                        # Regular key access
                        if isinstance(current, dict) and part in current:
                            current = current[part]
                        else:
                            return False
                
                # If we successfully traversed the path, it matches
                return True
            except Exception as e:
                logger.debug(f"Path matching failed for {useCase.useCaseId}: {e}")
        
        return False
    
    def _registerDefaultUseCases(self):
        """Register default use cases."""
        
        # Use Case 1: Section Content Generation
        # Returns JSON with "elements" array directly
        self.register(LoopingUseCase(
            useCaseId="section_content",
            jsonTemplate={"elements": []},
            detectionKeys=["elements"],
            detectionPath="",
            initialPromptBuilder=None,  # Will use default prompt builder
            continuationPromptBuilder=None,  # Will use default continuation builder
            accumulator=None,  # Direct return, no accumulation
            merger=None,
            continuationContextBuilder=None,  # Will use default continuation context
            resultBuilder=None,  # Return JSON directly
            supportsAccumulation=False,
            requiresExtraction=False
        ))
        
        # Use Case 2: Chapter Structure Generation
        # Returns JSON with "documents[0].chapters" structure
        self.register(LoopingUseCase(
            useCaseId="chapter_structure",
            jsonTemplate={"documents": [{"chapters": []}]},
            detectionKeys=["chapters"],
            detectionPath="documents[0].chapters",
            initialPromptBuilder=None,
            continuationPromptBuilder=None,
            accumulator=None,  # Direct return, no accumulation
            merger=None,
            continuationContextBuilder=None,
            resultBuilder=None,  # Return JSON directly
            supportsAccumulation=False,
            requiresExtraction=False
        ))
        
        # Use Case 3: Document Structure Generation
        # Returns JSON with "documents[0].sections" structure, requires extraction and accumulation
        self.register(LoopingUseCase(
            useCaseId="document_structure",
            jsonTemplate={"documents": [{"sections": []}]},
            detectionKeys=["sections"],
            detectionPath="documents[0].sections",
            initialPromptBuilder=None,
            continuationPromptBuilder=None,
            accumulator=None,  # Will use default accumulator
            merger=None,  # Will use default merger
            continuationContextBuilder=None,
            resultBuilder=None,  # Will use default result builder
            supportsAccumulation=True,
            requiresExtraction=True
        ))
        
        # Use Case 4: Code Structure Generation (NEW)
        self.register(LoopingUseCase(
            useCaseId="code_structure",
            jsonTemplate={
                "metadata": {
                    "language": "",
                    "projectType": "single_file|multi_file",
                    "projectName": ""
                },
                "files": [
                    {
                        "id": "",
                        "filename": "",
                        "fileType": "",
                        "dependencies": [],
                        "imports": [],
                        "functions": [],
                        "classes": []
                    }
                ]
            },
            detectionKeys=["files"],
            detectionPath="files",
            initialPromptBuilder=None,
            continuationPromptBuilder=None,
            accumulator=None,  # Direct return
            merger=None,
            continuationContextBuilder=None,
            resultBuilder=None,
            supportsAccumulation=False,
            requiresExtraction=False
        ))
        
        # Use Case 5: Code Content Generation (NEW)
        self.register(LoopingUseCase(
            useCaseId="code_content",
            jsonTemplate={"files": [{"content": "", "functions": []}]},
            detectionKeys=["content", "functions"],
            detectionPath="files[0].content",
            initialPromptBuilder=None,
            continuationPromptBuilder=None,
            accumulator=None,  # Will use default accumulator
            merger=None,  # Will use default merger
            continuationContextBuilder=None,
            resultBuilder=None,  # Will use default result builder
            supportsAccumulation=True,
            requiresExtraction=False
        ))
        
        # Use Case 6: Image Batch Generation (NEW)
        self.register(LoopingUseCase(
            useCaseId="image_batch",
            jsonTemplate={"images": []},
            detectionKeys=["images"],
            detectionPath="images",
            initialPromptBuilder=None,
            continuationPromptBuilder=None,
            accumulator=None,  # Direct return
            merger=None,
            continuationContextBuilder=None,
            resultBuilder=None,
            supportsAccumulation=False,
            requiresExtraction=False
        ))
        
        logger.info(f"Registered {len(self.useCases)} default looping use cases")