import asyncio
import aiohttp
import aiofiles
import tarfile
import shutil
from pathlib import Path
from typing import List, Set
import logging
from concurrent.futures import ThreadPoolExecutor
import os

# --- configuration ---
BASE_DIR = Path(__file__).parent.resolve() / Path("../../../..")
DATA_DIR = BASE_DIR / "data"
CSV_DIR = BASE_DIR / "data/csv"
EXTRACT_ADSB = BASE_DIR / "Code/cpp/extract-adsb"

# Performance tuning
MAX_CONCURRENT_DOWNLOADS = 2  # Limit concurrent HTTP requests
MAX_CONCURRENT_EXTRACTIONS = 4  # Limit concurrent tar extractions (CPU intensive)
MAX_CONCURRENT_DAYS = 3  # Limit concurrent day processing
# -----------------------

# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def pretty_path(path: os.PathLike | str) -> Path:
    """Return a path relative to the current working directory."""
    path = Path(path).resolve()
    return Path(os.path.relpath(path, Path.cwd()))

def pretty_url(url: str) -> str:
    """Return a shortened version of the URL for logging."""
    if len(url) <= 60:
        return url
    return f"{url[:30]}...{url[-30:]}"


async def download(session: aiohttp.ClientSession, url: str, dest: Path, semaphore: asyncio.Semaphore) -> None:
    """Download a file asynchronously with semaphore limiting and retry logic."""
    async with semaphore:
        dest.parent.mkdir(parents=True, exist_ok=True)
        logger.info(f"Downloading {pretty_url(url)} → {pretty_path(dest)}")
        
        max_retries = 3
        retry_delay = 10  # seconds
        timeout = 600  # seconds
        
        for attempt in range(max_retries + 1):
            try:
                async with session.get(url, timeout=aiohttp.ClientTimeout(total=timeout)) as response:
                    if response.status == 429:  # Too Many Requests
                        wait_time = int(response.headers.get('Retry-After', retry_delay * (attempt + 1)))
                        logger.warning(f"Rate limited for {pretty_url(url)}. Waiting {wait_time}s before retry {attempt + 1}/{max_retries}")
                        await asyncio.sleep(wait_time)
                        continue
                    
                    response.raise_for_status()
                    
                    async with aiofiles.open(dest, 'wb') as f:
                        async for chunk in response.content.iter_chunked(8192 * 4):
                            await f.write(chunk)
                            
                
                logger.info(f"Download completed: {pretty_path(dest)}")
                return  # Success, exit function
                
            except aiohttp.ClientResponseError as e:
                if e.status >= 500 and attempt < max_retries:  # Server errors are retryable
                    logger.warning(f"Server error {e.status} for {pretty_url(url)} on attempt {attempt + 1}/{max_retries}: {e}")
                    await asyncio.sleep(retry_delay * (2 ** attempt))  # Exponential backoff
                    continue
                else:
                    logger.error(f"HTTP error {e.status} for {pretty_url(url)}: {e}")
                    raise
                    
            except aiohttp.ClientConnectorError as e:
                if attempt < max_retries:
                    logger.warning(f"Connection error for {pretty_url(url)} on attempt {attempt + 1}/{max_retries}: {e}")
                    await asyncio.sleep(retry_delay * (2 ** attempt))
                    continue
                else:
                    logger.error(f"Connection failed for {pretty_url(url)}: {e}")
                    raise
                    
            except asyncio.TimeoutError as e:
                if attempt < max_retries:
                    logger.warning(f"Timeout for {pretty_url(url)} on attempt {attempt + 1}/{max_retries}: {e}")
                    await asyncio.sleep(retry_delay * (2 ** attempt))
                    continue
                else:
                    logger.error(f"Download timeout for {pretty_url(url)}: {e}")
                    raise
                    
            except Exception as e:
                if attempt < max_retries:
                    logger.warning(f"Unexpected error for {pretty_url(url)} on attempt {attempt + 1}/{max_retries}: {e}")
                    await asyncio.sleep(retry_delay * (2 ** attempt))
                    continue
                else:
                    logger.error(f"Download failed for {pretty_url(url)} after {max_retries} attempts: {e}")
                    raise
        
        # This should never be reached, but just in case
        raise Exception(f"Download failed for {pretty_url(url)} after {max_retries} attempts")

def get_name_from_url(url: str) -> str:
    """Extract filename from URL."""
    return url.split("/")[-1]


async def extract_adsb_data(cmd: List[str]) -> None:
    """Run extract-adsb binary asynchronously."""
    logger.info(f"Running: {' '.join(cmd)}")
    process = await asyncio.create_subprocess_exec(*cmd)
    await process.wait()
    if process.returncode != 0:
        raise Exception(f"Command failed with return code {process.returncode}: {' '.join(cmd)}")


def extract_tar_sync(tar_paths: List[Path], extract_dir: Path) -> Path:
    """Synchronous tar extraction (CPU-bound, run in thread pool)."""
    extract_dir.mkdir(parents=True, exist_ok=True)

    if len(tar_paths) == 1:
        tar_path = tar_paths[0]
    else:
        # Keep only multipart fragments like .aa, .ab, .ac ...
        parts = [p for p in tar_paths if len(p.suffix) == 3]

        # Derive correct base archive name: strip only the multipart suffix
        base = parts[0].with_suffix("")  # "foo.tar.aa" → "foo.tar"
        tar_path = extract_dir / base.name

        logger.info(f"Concatenating {len(parts)} parts → {pretty_path(tar_path)}")
        with open(tar_path, "wb") as outfile:
            for part in sorted(parts):
                with open(part, "rb") as infile:
                    shutil.copyfileobj(infile, outfile)

    logger.info(f"Extracting {pretty_path(tar_path)} → {pretty_path(extract_dir)}")
    with tarfile.open(tar_path) as tf:
        tf.extractall(path=extract_dir)

    return extract_dir


async def extract_tar(tar_paths: List[Path], extract_dir: Path, executor: ThreadPoolExecutor) -> Path:
    """Run tar extraction in thread pool to avoid blocking event loop."""
    loop = asyncio.get_event_loop()
    return await loop.run_in_executor(executor, extract_tar_sync, tar_paths, extract_dir)


async def move_csv_files(src_dir: Path, dest_dir: Path) -> None:
    """Move all CSV files from source to destination directory."""
    dest_dir.mkdir(parents=True, exist_ok=True)
    
    for p in src_dir.rglob("*.csv"):
        dest = dest_dir / p.name
        shutil.move(str(p), str(dest))


async def remove_dir(dir_path: Path) -> None:
    """Remove directory and all its contents."""
    if dir_path.exists():
        shutil.rmtree(dir_path)


async def remove_file(file_path: Path) -> None:
    """Remove a single file."""
    if file_path.exists():
        file_path.unlink()


async def download_day_parts(session: aiohttp.ClientSession, urls: List[str], 
                           download_semaphore: asyncio.Semaphore) -> List[Path]:
    """Download all parts for a single day concurrently."""
    tar_files = []
    download_tasks = []
    
    for url in urls:
        filename = DATA_DIR / get_name_from_url(url)
        tar_files.append(filename)
        download_tasks.append(
            download(session, url, filename, download_semaphore)
        )
    
    await asyncio.gather(*download_tasks)
    return tar_files


async def process_single_day(session: aiohttp.ClientSession, line: str, 
                           download_semaphore: asyncio.Semaphore,
                           extraction_executor: ThreadPoolExecutor) -> None:
    """Process a single day's data with concurrent downloads and parallel extraction."""
    parts = line.strip().split(",")
    if not parts or not parts[0]:
        return
    
    if len(parts) > 1:
        # Remove any part that is not of ending .tar.aa, .tar.ab, etc.
        parts = [p for p in parts if len(p.split(".")[-1]) == 2]
    
    logger.info(f"Processing day with {len(parts)} parts")
    
    try:
        # Download all parts for this day concurrently
        tar_files = await download_day_parts(session, parts, download_semaphore)
        
        # Extract the tar files
        extract_dir = tar_files[0].parent / tar_files[0].stem
        await extract_tar(tar_files, extract_dir, extraction_executor)
        
        # Process ADSB data
        traces_dir = extract_dir / "traces"
        if traces_dir.exists():
            await extract_adsb_data([str(EXTRACT_ADSB), str(traces_dir)])
        
        # Move CSV files
        csv_files_dir = CSV_DIR / extract_dir.name
        await move_csv_files(extract_dir, csv_files_dir)
        
        # Cleanup
        cleanup_tasks = [
            asyncio.create_task(remove_file(tar_file)) 
            for tar_file in tar_files 
            if tar_file.exists()
        ]
        cleanup_tasks.append(asyncio.create_task(remove_dir(extract_dir)))
        
        await asyncio.gather(*cleanup_tasks)
        logger.info(f"Completed processing day: {extract_dir.name}")
        
    except Exception as e:
        logger.error(f"Failed to process day: {e}")
        raise
    else:
        # If successful, write to ignore file
        await write_line_to_ignore(BASE_DIR / "Code/python/adsbpy/.ignore", line.strip())

async def write_line_to_ignore(file: Path, line: str) -> None:
    """Append a line to the ignore file asynchronously."""
    async with aiofiles.open(file, 'a') as f:
        await f.write(line + '\n')


async def process_days_concurrently(lines: List[str]) -> None:
    """Process multiple days concurrently with controlled parallelism."""
    # Create semaphores for controlling concurrency
    download_semaphore = asyncio.Semaphore(MAX_CONCURRENT_DOWNLOADS)
    day_semaphore = asyncio.Semaphore(MAX_CONCURRENT_DAYS)
    
    # Thread pool for CPU-bound tar extraction
    with ThreadPoolExecutor(max_workers=MAX_CONCURRENT_EXTRACTIONS) as extraction_executor:
        connector = aiohttp.TCPConnector(limit=MAX_CONCURRENT_DOWNLOADS)
        async with aiohttp.ClientSession(connector=connector) as session:
            
            async def process_day_with_semaphore(line: str) -> None:
                """Wrapper to limit concurrent day processing."""
                async with day_semaphore:
                    await process_single_day(session, line, download_semaphore, extraction_executor)
            
            # Process all days concurrently with limits
            tasks = [
                process_day_with_semaphore(line) 
                for line in lines 
                if line.strip()
            ]
            
            # Process in batches to avoid overwhelming the system
            batch_size = MAX_CONCURRENT_DAYS * 5
            for i in range(0, len(tasks), batch_size):
                batch = tasks[i:i + batch_size]
                logger.info(f"Processing batch {i//batch_size + 1}/{(len(tasks)-1)//batch_size + 1}")
                await asyncio.gather(*batch, return_exceptions=True)


async def fetch_lines(lines: List[str]) -> None:
    """Process multiple lines (days) concurrently."""
    await process_days_concurrently(lines)


async def fetch_file(file: Path) -> None:
    """Fetch and process URLs from a file."""
    async with aiofiles.open(file, 'r') as f:
        lines = await f.readlines()
    
    logger.info(f"Processing {len(lines)} days")
    await fetch_lines(lines)


async def main() -> None:
    """Main entry point if running as a script."""
    import sys
    
    if len(sys.argv) > 1:
        input_file = Path(sys.argv[1])
    else:
        input_file = Path("urls.txt")
    
    if input_file.exists():
        logger.info(f"Starting processing of {input_file}")
        await fetch_file(input_file)
        logger.info("Processing completed successfully")
    else:
        logger.error(f"Input file {input_file} not found")
        sys.exit(1)


if __name__ == "__main__":
    asyncio.run(main())