from prefect import flow, task
import asyncio
from typing import List, Dict
import xml.etree.ElementTree as ET
from google.cloud import storage
import io
from datetime import datetime

# Simulating GCS file content
SAMPLE_XML = """
<filing>
    <header>
        <filing_id>{id}</filing_id>
        <timestamp>{timestamp}</timestamp>
    </header>
    <data>
        <field1>value1</field1>
        <field2>value2</field2>
        <!-- Imagine more fields here -->
    </data>
</filing>
"""

# Simulating GCS operations
async def download_from_gcs(file_path: str) -> str:
    """Simulate downloading a file from GCS"""
    # In real code, this would actually download from GCS
    await asyncio.sleep(0.1)  # Simulate network delay
    return SAMPLE_XML.format(
        id=file_path,
        timestamp=datetime.now().isoformat()
    )

async def process_single_file(file_path: str, semaphore: asyncio.Semaphore) -> Dict:
    """Process a single XML file with semaphore control"""
    async with semaphore:
        content = await download_from_gcs(file_path)
        # Parse XML and extract fields
        root = ET.fromstring(content)
        
        # Extract specific fields
        result = {
            'filing_id': root.find('.//filing_id').text,
            'timestamp': root.find('.//timestamp').text,
            'field1': root.find('.//field1').text,
            'field2': root.find('.//field2').text
        }
        
        return result

# Version 1: Without @task decorator
async def process_batch_no_decorator(file_paths: List[str]) -> List[Dict]:
    """Process a batch of files without the @task decorator"""
    semaphore = asyncio.Semaphore(20)  # Limit concurrent downloads
    results = await asyncio.gather(
        *[process_single_file(path, semaphore) for path in file_paths]
    )
    return results

# Version 2: With @task decorator
@task(task_run_name="process_batch", persist_result=False)
async def process_batch_with_decorator(file_paths: List[str]) -> List[Dict]:
    """Process a batch of files with the @task decorator"""
    semaphore = asyncio.Semaphore(20)  # Limit concurrent downloads
    results = await asyncio.gather(
        *[process_single_file(path, semaphore) for path in file_paths]
    )
    return results

@flow
async def main_flow(use_task_decorator: bool = True):
    """Main flow that processes multiple batches"""
    # Simulate a large number of files
    all_files = [f"file_{i}.xml" for i in range(100000)]
    batch_size = 500
    
    results = []
    for i in range(0, len(all_files), batch_size):
        batch = all_files[i:i + batch_size]
        if use_task_decorator:
            batch_results = await process_batch_with_decorator(batch)
        else:
            batch_results = await process_batch_no_decorator(batch)
            
        # Process results (in real code, might write to storage here)
        results.extend(batch_results)
        
        # Clear batch results to free memory
        batch_results = None
    
    return len(results)

# Run the flow
if __name__ == "__main__":
    import sys
    use_decorator = len(sys.argv) > 1 and sys.argv[1].lower() == "true"
    print(f"Running with task decorator: {use_decorator}")
    
    asyncio.run(main_flow(use_decorator))