Check sources coverage

import asyncio
import csv
import logging
import os
from typing import Dict, List, Optional

from newscatcher import AsyncNewscatcherApi

INPUT_FILE = "source_urls.csv"
OUTPUT_FILE = "uncovered_sources.csv"
BATCH_SIZE = 1000
MAX_CONCURRENT = 5

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)

client = AsyncNewscatcherApi(api_key=os.environ["NEWSCATCHER_API_KEY"])


async def check_batch(
    semaphore: asyncio.Semaphore,
    batch: List[str],
) -> Optional[object]:
    async with semaphore:
        try:
            return await client.sources.post(
                source_url=batch,
                include_additional_info=True,
            )
        except Exception as e:
            logging.error(f"Batch failed: {e}")
            return None


async def check_coverage(source_urls: List[str]) -> Dict[str, bool]:
    coverage: Dict[str, bool] = {}
    semaphore = asyncio.Semaphore(MAX_CONCURRENT)

    batches = [
        source_urls[i : i + BATCH_SIZE]
        for i in range(0, len(source_urls), BATCH_SIZE)
    ]

    results = await asyncio.gather(
        *[check_batch(semaphore, batch) for batch in batches],
        return_exceptions=True,
    )

    for i, result in enumerate(results):
        if isinstance(result, Exception) or result is None:
            logging.error(f"Skipping failed batch {i + 1}")
            continue
        for source in result.sources or []:
            coverage[source.domain_url] = True
        for url in batches[i]:
            if url not in coverage:
                coverage[url] = False

    return coverage


def read_sources(file_path: str) -> List[str]:
    with open(file_path) as f:
        reader = csv.reader(f)
        sources = [row[0] for row in reader if row]
    logging.info(f"Loaded {len(sources)} sources from {file_path}")
    return sources


def write_uncovered(uncovered: List[str], file_path: str) -> None:
    with open(file_path, "w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(["Uncovered Source URL"])
        for url in uncovered:
            writer.writerow([url])
    logging.info(f"Wrote {len(uncovered)} uncovered sources to {file_path}")


def main() -> None:
    source_urls = read_sources(INPUT_FILE)
    coverage = asyncio.run(check_coverage(source_urls))
    uncovered = [url for url, covered in coverage.items() if not covered]
    write_uncovered(uncovered, OUTPUT_FILE)
    logging.info(
        f"Total: {len(source_urls)} | "
        f"Covered: {len(source_urls) - len(uncovered)} | "
        f"Uncovered: {len(uncovered)}"
    )


if __name__ == "__main__":
    main()