import logging
import os
import sys
from datetime import datetime
from io import BytesIO
from pathlib import Path

import niquests
from extract_core import (
    end_run,
    landing_path,
    normalize_etag,
    open_state_db,
    start_run,
    write_bytes_atomic,
)

from .normalize import normalize_zipped_csv

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 = 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) -> int:
    """Extract PSD file to local year/month subdirectory.

    Returns bytes_written (0 if the file already existed and was skipped).
    """
    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 0
    elif response.status_code != 200:
        logger.error(f"Status code not ok, STATUS={response.status_code}")
        return 0

    raw_etag = response.headers.get("etag", "")
    assert raw_etag, "USDA response missing etag header"
    etag = normalize_etag(raw_etag)

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

    response = http_session.get(url, timeout=HTTP_TIMEOUT_SECONDS)
    logger.info(f"Storing file to {local_file}")
    normalized_content = normalize_zipped_csv(BytesIO(response.content))
    bytes_written = write_bytes_atomic(local_file, normalized_content.read())
    logger.info("Download complete")
    return bytes_written


def extract_psd_dataset():
    conn = open_state_db(LANDING_DIR)
    run_id = start_run(conn, "psdonline")
    files_written = 0
    files_skipped = 0
    bytes_written = 0
    cursor_value = None
    try:
        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}")
                    result = extract_psd_file(url=url, year=year, month=month, http_session=session)
                    if result > 0:
                        files_written = 1
                        bytes_written = result
                    else:
                        files_skipped = 1
                    cursor_value = f"{year}-{month:02d}"
                    break
                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}")
            else:
                logger.error("Could not find any available data in the last 4 months")

        end_run(
            conn, run_id, status="success",
            files_written=files_written, files_skipped=files_skipped,
            bytes_written=bytes_written, cursor_value=cursor_value,
        )
    except Exception as e:
        end_run(conn, run_id, status="failed", error_message=str(e))
        raise
    finally:
        conn.close()


if __name__ == "__main__":
    extract_psd_dataset()
