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)

    @staticmethod
    def get_ai_selected_high_quality_datasets() -> List[str]:
        """
        To remove personal bias, Google gemini was asked to select ten high quality galaxies
        for testing.
        Working with these galaxies reduces the MAPE from ~23% to ~13%.
        | Galaxy Name |     Inclination (i)  |   Quality Note   |   Why it's good for DG |
        | ----------- | -------------------- | ---------------- | ---------------------- |
        | NGC 2403    |    63∘               | High Resolution   |  Exceptional precision; the ''hydrogen standard.'' |
        | NGC 3198    |    72∘               | Extended RC       |  Very famous ''classic'' for testing flat rotation curves. |
        | NGC 2903    |    65∘               | Multi-Component   |  Clear disk and gas profiles; minimizes Υ ambiguity. |
        | NGC 5055    |    59∘               | Large Mass Range  |  Massive galaxy with a well-defined outer plateau. |
        | NGC 7331    |    76∘               | HSB Standard      |  Excellent for testing the transition into the Machian regime. |
        | DDO 154     |    66∘               | Gas Dominated     |  Dark matter dominated; perfect for a0​ testing. |
        | NGC 3741    |    60∘               | Hyper-Extended    |  Traces the potential out to extreme radii (Machian dominance). |
        | NGC 6946    |    33∘               | Quality Fit       |  While lower, its quality is so high it is often prioritized. |
        | UGC 128     |    54∘               | LSB Prototype     |  Classic Low Surface Brightness galaxy for gravity tests. |
        | NGC 3953    |    62∘               | Stable Disk       |  Very low non-circular motions (cleaner data). |
        """
        return [
            "NGC2403",
            "NGC3198",
            "NGC2903",
            "NGC5055",
            "NGC7331",
            "DDO154",
            "NGC3741",
            "NGC6946",
            "UGC00128",
            "NGC3953",
        ]
