import json
import logging
import os
import tarfile
import threading
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Dict, Iterable, List, Optional

import requests


AZURE_SAS_ENDPOINT = "https://planetarycomputer.microsoft.com/api/sas/v1/token/noaagefs/gefs"
FORECAST_HOURS_GEFS: List[int] = list(range(0, 387, 3))
FORECAST_HOURS_GFS: List[int] = list(range(0, 387, 3))
STATE_FILENAME = "state.json"


logger = logging.getLogger("gefs_wave")
if not logger.handlers:
    handler = logging.StreamHandler()
    handler.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(message)s"))
    logger.addHandler(handler)
logger.setLevel(logging.INFO)
logger.propagate = False


@dataclass
class Endpoint:
    base_url: str
    mode: str = "dated"  # "dated" or "latest"
    needs_sas: bool = False
    label: Optional[str] = None


@dataclass
class WaveProduct:
    key: str
    label: str
    directory_template: str
    filename_template: str
    endpoints: List[Endpoint]
    forecast_hours: List[int]
    max_days_back: int


PRODUCTS: List[WaveProduct] = [
    WaveProduct(
        key="gefswave_c00",
        label="GEFS Wave Control (c00)",
        directory_template="gefs.{date}/{cycle}/wave/gridded",
        filename_template="gefswave.t{cycle}z.c00.global.0p25.f{hour:03d}.grib2",
        endpoints=[
            Endpoint("https://noaagefs.blob.core.windows.net/gefs", needs_sas=True, label="Azure GEFS mirror"),
            Endpoint("https://nomads.ncep.noaa.gov/pub/data/nccf/com/gens/prod", label="NOMADS GEFS"),
        ],
        forecast_hours=FORECAST_HOURS_GEFS,
        max_days_back=9,
    ),
    WaveProduct(
        key="gfswave_global",
        label="GFS Wave Global 0.25°",
        directory_template="gfs.{date}/{cycle}/wave/gridded",
        filename_template="gfswave.t{cycle}z.global.0p25.f{hour:03d}.grib2",
        endpoints=[
            Endpoint("https://nomads.ncep.noaa.gov/pub/data/nccf/com/wave/prod", label="NOMADS GFS"),
            Endpoint("https://polar.ncep.noaa.gov/waves/latest_run", mode="latest", label="NOAA Polar latest_run"),
        ],
        forecast_hours=FORECAST_HOURS_GFS,
        max_days_back=5,
    ),
]


@dataclass
class Cycle:
    date: str  # YYYYMMDD
    cycle: str  # HH

    @property
    def label(self) -> str:
        return f"{self.date} {self.cycle}Z"


class WaveDownloadError(Exception):
    """Raised when a download or discovery step fails."""


class RetryableAuthError(WaveDownloadError):
    """Raised when an auth token should be refreshed and retried."""


class WaveDownloader:
    """Handles discovery and download of wave datasets."""

    def __init__(self, data_root: Path, session: Optional[requests.Session] = None) -> None:
        self.data_root = data_root
        self.data_root.mkdir(parents=True, exist_ok=True)
        self.session = session or requests.Session()
        self._azure_sas_token: Optional[str] = None
        self._azure_token_checked = False

    def ensure_latest_cycle(self) -> Dict:
        """Ensure the latest available dataset (GEFS first, then GFS) is cached locally."""
        errors: List[str] = []
        for product in PRODUCTS:
            try:
                return self._ensure_product(product)
            except WaveDownloadError as exc:
                logger.warning("Failed to refresh %s: %s", product.label, exc)
                errors.append(f"{product.label}: {exc}")

        raise WaveDownloadError("All configured products failed. " + " | ".join(errors))

    def _ensure_product(self, product: WaveProduct) -> Dict:
        logger.info("Ensuring latest %s dataset in %s", product.label, self.data_root)
        cycle = self._find_latest_cycle(product)
        if not cycle:
            raise WaveDownloadError("Could not locate a recent cycle.")

        state = self._read_state()
        if (
            state
            and state.get("product") == product.key
            and state.get("cycle") == cycle.label
            and self._cycle_complete(product, cycle)
        ):
            logger.info("Cycle %s already cached for %s.", cycle.label, product.label)
            return state

        logger.info("Downloading %s cycle %s", product.label, cycle.label)
        files = self._download_cycle(product, cycle)
        tarball_path = self._build_tarball(product, cycle, files)
        state = {
            "product": product.key,
            "product_label": product.label,
            "cycle": cycle.label,
            "files": [str(path.relative_to(self.data_root)) for path in files],
            "tarball": str(tarball_path.relative_to(self.data_root)),
            "updated_at": datetime.now(timezone.utc).isoformat(),
        }
        self._write_state(state)
        return state

    def current_state(self) -> Optional[Dict]:
        return self._read_state()

    def _find_latest_cycle(self, product: WaveProduct) -> Optional[Cycle]:
        now = datetime.now(timezone.utc)
        candidate_days = [(now - timedelta(days=days)).strftime("%Y%m%d") for days in range(product.max_days_back + 1)]
        cycles = ["18", "12", "06", "00"]

        for day in candidate_days:
            for cycle in cycles:
                logger.info("Probing %s cycle %s %sz", product.label, day, cycle)
                if self._cycle_available(product, day, cycle):
                    logger.info("Found %s cycle %s %sz.", product.label, day, cycle)
                    return Cycle(date=day, cycle=cycle)
        logger.warning("No %s cycle found in the last %d days.", product.label, product.max_days_back)
        return None

    def _cycle_available(self, product: WaveProduct, day: str, cycle_hour: str) -> bool:
        for endpoint in product.endpoints:
            if self._probe_endpoint(product, endpoint, day, cycle_hour):
                return True
        return False

    def _probe_endpoint(self, product: WaveProduct, endpoint: Endpoint, day: str, cycle_hour: str) -> bool:
        url = self._build_url(product, endpoint, day, cycle_hour, hour=0)
        if not url:
            return False

        try:
            response = self.session.head(url, timeout=20)
            if response.ok:
                logger.info("Cycle %s %sz available via %s", day, cycle_hour, endpoint.label or endpoint.base_url)
                return True

            if response.status_code in (401, 403) and endpoint.needs_sas:
                logger.warning(
                    "Auth error probing %s on %s (HTTP %s). Refreshing credentials.",
                    product.label,
                    endpoint.label or endpoint.base_url,
                    response.status_code,
                )
                self._invalidate_azure_token()
                url = self._build_url(product, endpoint, day, cycle_hour, hour=0)
                if not url:
                    return False
                response = self.session.head(url, timeout=20)
                if response.ok:
                    return True

            if response.status_code not in (404, 400):
                range_headers = {"Range": "bytes=0-0"}
                with self.session.get(url, headers=range_headers, timeout=20, stream=True) as stream_resp:
                    if stream_resp.status_code in (200, 206):
                        logger.info(
                            "Cycle %s %sz confirmed via range request on %s",
                            day,
                            cycle_hour,
                            endpoint.label or endpoint.base_url,
                        )
                        return True
        except requests.RequestException as exc:
            logger.debug(
                "Error probing %s cycle %s %sz on %s: %s",
                product.label,
                day,
                cycle_hour,
                endpoint.label or endpoint.base_url,
                exc,
            )
        return False

    def _download_cycle(self, product: WaveProduct, cycle: Cycle) -> List[Path]:
        files: List[Path] = []
        target_dir = self.data_root / product.key / cycle.date / cycle.cycle
        target_dir.mkdir(parents=True, exist_ok=True)

        for hour in product.forecast_hours:
            filename = product.filename_template.format(cycle=cycle.cycle, hour=hour)
            destination = target_dir / filename
            if destination.exists() and destination.stat().st_size > 0:
                logger.debug("File %s already exists; skipping.", destination)
                files.append(destination)
                continue
            logger.info("Downloading %s", filename)
            self._download_with_fallback(product, cycle, filename, destination)
            files.append(destination)
        return files

    def _build_tarball(self, product: WaveProduct, cycle: Cycle, files: Iterable[Path]) -> Path:
        tarball_path = self.data_root / f"{product.key}-latest.tar.gz"
        with tarfile.open(tarball_path, "w:gz") as tar:
            for file_path in files:
                tar.add(file_path, arcname=file_path.relative_to(self.data_root))
        return tarball_path

    def _download_with_fallback(self, product: WaveProduct, cycle: Cycle, filename: str, destination: Path) -> None:
        errors = []
        for endpoint in product.endpoints:
            try:
                self._stream_to_file(product, endpoint, cycle, filename, destination)
                return
            except RetryableAuthError:
                logger.info("Retrying %s download on %s after refreshing credentials.", filename, endpoint.base_url)
                self._stream_to_file(product, endpoint, cycle, filename, destination)
                return
            except WaveDownloadError as exc:
                errors.append(str(exc))
                logger.warning("Failed to download %s from %s (%s).", filename, endpoint.base_url, exc)
        raise WaveDownloadError(f"All download attempts failed for {filename}: {'; '.join(errors)}")

    def _stream_to_file(
        self,
        product: WaveProduct,
        endpoint: Endpoint,
        cycle: Cycle,
        filename: str,
        destination: Path,
    ) -> None:
        url = self._build_url(product, endpoint, cycle.date, cycle.cycle, filename=filename)
        if not url:
            raise WaveDownloadError(f"No accessible URL for endpoint {endpoint.base_url}")

        tmp_path = destination.with_suffix(destination.suffix + ".part")
        try:
            with self.session.get(url, stream=True, timeout=120) as response:
                if response.status_code in (401, 403) and endpoint.needs_sas:
                    self._invalidate_azure_token()
                    raise RetryableAuthError(f"Auth failure downloading {filename} from {endpoint.base_url}")
                response.raise_for_status()
                with tmp_path.open("wb") as handle:
                    for chunk in response.iter_content(chunk_size=1 << 20):
                        if not chunk:
                            continue
                        handle.write(chunk)
            tmp_path.rename(destination)
        except RetryableAuthError:
            if tmp_path.exists():
                tmp_path.unlink()
            raise
        except requests.RequestException as exc:
            if tmp_path.exists():
                tmp_path.unlink()
            logger.error("Failed to download %s: %s", url, exc)
            raise WaveDownloadError(f"Failed to download {url}: {exc}") from exc

    def _build_url(
        self,
        product: WaveProduct,
        endpoint: Endpoint,
        day: str,
        cycle_hour: str,
        hour: Optional[int] = None,
        filename: Optional[str] = None,
    ) -> Optional[str]:
        if filename is None:
            if hour is None:
                raise ValueError("Either hour or filename must be provided.")
            filename = product.filename_template.format(cycle=cycle_hour, hour=hour)

        base = endpoint.base_url.rstrip("/")
        if endpoint.mode == "latest":
            url = f"{base}/{filename}"
        else:
            directory = product.directory_template.format(date=day, cycle=cycle_hour)
            url = f"{base}/{directory}/{filename}"

        if endpoint.needs_sas:
            token = self._get_azure_sas_token()
            if not token:
                return None
            token = token if token.startswith("?") else f"?{token}"
            return f"{url}{token}"
        return url

    def _get_azure_sas_token(self) -> Optional[str]:
        if self._azure_sas_token:
            return self._azure_sas_token

        env_token = os.environ.get("GEFS_AZURE_SAS_TOKEN")
        if env_token:
            self._azure_sas_token = env_token.strip()
            logger.info("Using Azure SAS token from environment variable.")
            return self._azure_sas_token

        if self._azure_token_checked:
            return None

        self._azure_token_checked = True
        try:
            response = self.session.get(AZURE_SAS_ENDPOINT, timeout=20)
            response.raise_for_status()
            token = response.json().get("token")
            if token:
                self._azure_sas_token = token.strip()
                logger.info("Fetched Azure SAS token from Planetary Computer API.")
            else:
                logger.warning("Azure SAS endpoint response did not include a token.")
        except requests.RequestException as exc:
            logger.warning("Failed to fetch Azure SAS token: %s", exc)
        return self._azure_sas_token

    def _invalidate_azure_token(self) -> None:
        if self._azure_sas_token:
            logger.info("Invalidating cached Azure SAS token.")
        self._azure_sas_token = None
        self._azure_token_checked = False

    def _cycle_complete(self, product: WaveProduct, cycle: Cycle) -> bool:
        target_dir = self.data_root / product.key / cycle.date / cycle.cycle
        if not target_dir.exists():
            return False
        expected = {
            product.filename_template.format(cycle=cycle.cycle, hour=hour) for hour in product.forecast_hours
        }
        existing = {path.name for path in target_dir.glob("*.grib2") if path.stat().st_size > 0}
        return expected.issubset(existing)

    def _read_state(self) -> Optional[Dict]:
        path = self.data_root / STATE_FILENAME
        if not path.exists():
            return None
        try:
            return json.loads(path.read_text())
        except json.JSONDecodeError:
            return None

    def _write_state(self, state: Dict) -> None:
        path = self.data_root / STATE_FILENAME
        path.write_text(json.dumps(state, indent=2))


class WaveDownloadManager:
    """Simple threaded wrapper to keep download work off the request thread."""

    def __init__(self, data_root: Path) -> None:
        self.downloader = WaveDownloader(data_root)
        self._lock = threading.Lock()
        self._worker: Optional[threading.Thread] = None
        self._status: Dict = {"status": "idle"}

    def trigger_refresh(self) -> None:
        with self._lock:
            if self._worker and self._worker.is_alive():
                logger.info("Refresh already in progress; ignoring new trigger.")
                return
            logger.info("Starting wave refresh thread.")
            self._status = {"status": "running"}
            self._worker = threading.Thread(target=self._run_refresh, daemon=True)
            self._worker.start()

    def status(self) -> Dict:
        with self._lock:
            status = dict(self._status)
        state = self.downloader.current_state()
        if state:
            status["latest_state"] = state
        return status

    def _run_refresh(self) -> None:
        try:
            result = self.downloader.ensure_latest_cycle()
            with self._lock:
                self._status = {"status": "ready", "latest_state": result}
            logger.info("Wave refresh completed successfully.")
        except Exception as exc:
            with self._lock:
                self._status = {"status": "error", "message": str(exc)}
            logger.exception("Wave refresh failed: %s", exc)