from dataclasses import dataclass
from pathlib import Path
from typing import List, Optional, Union
import os
import random
import requests
import zipfile
import platformdirs

DATA_URL = "https://zenodo.org/record/16284118/files/Rotmod_LTG.zip"
LOCAL_DIR = "sparc_cache"


@dataclass
class SparcHandler:
    output_dir: Union[Path, None] = None
    have_logged_cache: bool = False

    def __post_init__(self) -> None:
        if self.output_dir is None:
            self.output_dir = Path(platformdirs.user_data_dir())

    def initialize_sparc_data(self) -> zipfile.ZipFile:
        """Handles 404/429 errors by caching data locally."""
        if self.output_dir is None:
            print("Cannot continue without a cache directory.")
            exit(1)
        if not os.path.exists(self.output_dir / "divergent_gravity"):
            os.makedirs(self.output_dir / "divergent_gravity")

        zip_file = os.path.join(self.output_dir, "divergent_gravity", "Rotmod_LTG.zip")
        if not os.path.exists(zip_file):
            print("Attempting to fetch SPARC database from Zenodo mirror...")
            # Browser-like headers to prevent 403/429
            headers = {"User-Agent": "Mozilla/5.0"}
            try:
                r = requests.get(DATA_URL, headers=headers, stream=True)
                r.raise_for_status()
                with open(zip_file, "wb") as f:
                    f.write(r.content)
                print("Download successful.")
            except Exception as e:
                print(f"\nCRITICAL: Automatic download failed ({e}).")
                print(f"Please manually download the 'Rotmod_LTG.zip' from:")
                print(f"https://zenodo.org/records/16284118")
                print(
                    f"And place it in the folder: {os.path.abspath(self.output_dir / 'divergent_gravity')}"
                )
                exit(1)
        else:
            if not self.have_logged_cache:
                print(f"Using cached data from {zip_file}")
                self.have_logged_cache = True

        return zipfile.ZipFile(zip_file)

    def catalog_filename_to_name(self, catalog_filename: str) -> str:
        return catalog_filename.split("_")[0]

    def all_catalog_names(self):
        """Returns a list of catalog names in the SPARC database."""
        zf = self.initialize_sparc_data()
        return [self.catalog_filename_to_name(s) for s in zf.namelist()]

    def get_random_datasets(self, n_datasets: int = 5) -> List[str]:
        names = self.all_catalog_names()
        return random.sample(names, n_datasets)
