from .normalize import normalize_zipped_csv
import logging
import os
import pathlib
import sys
from datetime import datetime
from io import BytesIO

import niquests

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    handlers=[logging.StreamHandler(sys.stdout)],
)
logger = logging.getLogger("PSDOnline Extractor")

LANDING_DIR = pathlib.Path(os.getenv("LANDING_DIR", "data/landing"))
LANDING_DIR.mkdir(parents=True, exist_ok=True)
logger.info(f"Landing dir: {LANDING_DIR}")

PSD_HISTORICAL_URL = "https://apps.fas.usda.gov/psdonline/downloads/archives/{year}/{month:02d}/psd_alldata_csv.zip"
FIRST_YEAR = 2006
FIRST_MONTH = 8

HTTP_TIMEOUT_SECONDS = 60


def extract_psd_file(url: str, year: int, month: int, http_session: niquests.Session):
    """Extract PSD file to local year/month subdirectory."""
    logger.info(f"Requesting file {url} ...")

    response = http_session.head(url, timeout=HTTP_TIMEOUT_SECONDS)
    if response.status_code == 404:
        logger.error("File doesn't exist on server, received status code 404 Not Found")
        return
    elif response.status_code != 200:
        logger.error(f"Status code not ok, STATUS={response.status_code}")
        return

    etag = response.headers.get("etag", "").replace('"', "").replace(":", "_")
    assert etag, "USDA response missing etag header"

    extract_to_path = LANDING_DIR / "psd" / str(year) / f"{month:02d}"
    local_file = extract_to_path / f"{etag}.csv.gzip"
    if local_file.exists():
        logger.info(f"File {etag}.csv.gzip already exists locally, skipping")
        return

    response = http_session.get(url, timeout=HTTP_TIMEOUT_SECONDS)
    logger.info(f"Storing file to {local_file}")
    extract_to_path.mkdir(parents=True, exist_ok=True)
    normalized_content = normalize_zipped_csv(BytesIO(response.content))
    local_file.write_bytes(normalized_content.read())
    assert local_file.exists(), f"File was not written: {local_file}"
    logger.info("Download complete")


def extract_psd_dataset():
    today = datetime.now()

    with niquests.Session() as session:
        for months_back in range(4):
            year = today.year
            month = today.month - months_back
            while month < 1:
                month += 12
                year -= 1

            url = PSD_HISTORICAL_URL.format(year=year, month=month)
            logger.info(f"Trying {year}-{month:02d}...")

            response = session.head(url, timeout=HTTP_TIMEOUT_SECONDS)
            if response.status_code == 200:
                logger.info(f"Found latest data at {year}-{month:02d}")
                extract_psd_file(url=url, year=year, month=month, http_session=session)
                return
            elif response.status_code == 404:
                logger.info(f"Month {year}-{month:02d} not found, trying earlier...")
            else:
                logger.warning(f"Unexpected status code {response.status_code} for {year}-{month:02d}")

        logger.error("Could not find any available data in the last 4 months")


if __name__ == "__main__":
    extract_psd_dataset()
