"""USDA data service for Edge API SDK."""

from typing import Optional, Union
from datetime import date, datetime
import pandas as pd


class USDAService:
    """Service for accessing USDA data."""

    def __init__(self, client):
        """Initialize with parent client."""
        self.client = client
        self.source_name = "USDA"

    def get_available_series(self) -> pd.DataFrame:
        """
        Get list of available USDA time series.

        Returns:
            DataFrame with available USDA series metadata
        """
        return self.client.time_series.get_metadata(
            source_name=self.source_name,
            limit=500
        )

    def get_ams_data(
        self,
        series_code: str,
        start_date: Optional[Union[str, date, datetime]] = None,
        end_date: Optional[Union[str, date, datetime]] = None,
        limit: Optional[int] = None,
    ) -> pd.DataFrame:
        """
        Get AMS (Agricultural Marketing Service) data.

        Args:
            series_code: AMS series code
            start_date: Start date for data
            end_date: End date for data
            limit: Maximum records to return

        Returns:
            DataFrame with AMS data
        """
        return self.client.time_series.get_data(
            series_code=series_code,
            source_name=self.source_name,
            start_date=start_date,
            end_date=end_date,
            limit=limit,
        )

    def get_beef_prices(
        self,
        cut_type: Optional[str] = None,
        start_date: Optional[Union[str, date, datetime]] = None,
        end_date: Optional[Union[str, date, datetime]] = None,
        limit: Optional[int] = None,
    ) -> pd.DataFrame:
        """
        Get beef prices from USDA AMS.

        Args:
            cut_type: Specific cut type (e.g., 'ribeye', 'chuck', 'round')
            start_date: Start date for data
            end_date: End date for data
            limit: Maximum records to return

        Returns:
            DataFrame with beef price data
        """
        # Get metadata for beef series
        series_filter = "BEEF_"
        if cut_type:
            series_filter += cut_type.upper()

        # Find matching series
        metadata = self.client.time_series.search(
            search_term=series_filter,
            source_name=self.source_name,
            limit=100
        )

        if metadata.empty:
            return pd.DataFrame()

        # Get data for all matching series
        all_data = []
        for _, series in metadata.iterrows():
            series_data = self.client.time_series.get_data(
                series_code=series["series_code"],
                source_name=self.source_name,
                start_date=start_date,
                end_date=end_date,
                limit=limit,
            )
            if not series_data.empty:
                series_data["description"] = series.get("series_description", "")
                all_data.append(series_data)

        if all_data:
            return pd.concat(all_data, ignore_index=True)
        return pd.DataFrame()

    def get_pork_prices(
        self,
        cut_type: Optional[str] = None,
        start_date: Optional[Union[str, date, datetime]] = None,
        end_date: Optional[Union[str, date, datetime]] = None,
        limit: Optional[int] = None,
    ) -> pd.DataFrame:
        """
        Get pork prices from USDA AMS.

        Args:
            cut_type: Specific cut type (e.g., 'belly', 'ham', 'loin')
            start_date: Start date for data
            end_date: End date for data
            limit: Maximum records to return

        Returns:
            DataFrame with pork price data
        """
        series_filter = "PORK_"
        if cut_type:
            series_filter += cut_type.upper()

        metadata = self.client.time_series.search(
            search_term=series_filter,
            source_name=self.source_name,
            limit=100
        )

        if metadata.empty:
            return pd.DataFrame()

        all_data = []
        for _, series in metadata.iterrows():
            series_data = self.client.time_series.get_data(
                series_code=series["series_code"],
                source_name=self.source_name,
                start_date=start_date,
                end_date=end_date,
                limit=limit,
            )
            if not series_data.empty:
                series_data["description"] = series.get("series_description", "")
                all_data.append(series_data)

        if all_data:
            return pd.concat(all_data, ignore_index=True)
        return pd.DataFrame()

    def get_cutout_values(
        self,
        commodity: str = "beef",
        start_date: Optional[Union[str, date, datetime]] = None,
        end_date: Optional[Union[str, date, datetime]] = None,
    ) -> pd.DataFrame:
        """
        Get cutout values for beef or pork.

        Args:
            commodity: 'beef' or 'pork'
            start_date: Start date
            end_date: End date

        Returns:
            DataFrame with cutout values
        """
        series_code = f"{commodity.upper()}_CUTOUT"

        return self.client.time_series.get_data(
            series_code=series_code,
            source_name=self.source_name,
            start_date=start_date,
            end_date=end_date,
        )

    def get_nass_data(
        self,
        series_code: str,
        start_date: Optional[Union[str, date, datetime]] = None,
        end_date: Optional[Union[str, date, datetime]] = None,
    ) -> pd.DataFrame:
        """
        Get NASS (National Agricultural Statistics Service) data.

        Args:
            series_code: NASS series code
            start_date: Start date
            end_date: End date

        Returns:
            DataFrame with NASS data
        """
        return self.client.time_series.get_data(
            series_code=series_code,
            source_name="NASS",
            start_date=start_date,
            end_date=end_date,
        )