import os
import datetime
from google.cloud import bigquery

class BillingAgent:
    def __init__(self):
        self.project_id = os.environ.get("GOOGLE_CLOUD_PROJECT")
        if not self.project_id:
            raise ValueError("GOOGLE_CLOUD_PROJECT environment variable not set.")
        
        self.billing_table = os.environ.get("BILLING_TABLE")
        if not self.billing_table:
            raise ValueError("BILLING_TABLE environment variable not set.")

        self.bq_client = bigquery.Client(project=self.project_id)

    def get_summary(self):
        """
        Retrieves the billing summary for the last 30 days.
        """
        query = f"""
            SELECT
                service.description as service,
                SUM(cost) as total_cost
            FROM `{self.billing_table}`
            WHERE _PARTITIONTIME >= TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL 30 DAY)
            GROUP BY 1
            ORDER BY total_cost DESC
            LIMIT 10
        """
        query_job = self.bq_client.query(query)
        results = query_job.result()
        
        summary = []
        for row in results:
            summary.append({"service": row.service, "total_cost": row.total_cost})
            
        return {"summary": summary}

    def get_forecast(self):
        """
        Retrieves a billing forecast based on the last 7 days.
        """
        # Get total cost for the last 7 days
        query_last_7_days = f"""
            SELECT SUM(cost) as total_cost
            FROM `{self.billing_table}`
            WHERE _PARTITIONTIME >= TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL 7 DAY)
        """
        query_job = self.bq_client.query(query_last_7_days)
        results = query_job.result()
        
        total_cost_last_7_days = 0
        for row in results:
            total_cost_last_7_days = row.total_cost or 0

        daily_average = total_cost_last_7_days / 7

        # Calculate remaining days in the month
        today = datetime.date.today()
        last_day_of_month = datetime.date(today.year, today.month, 1) + datetime.timedelta(days=32)
        last_day_of_month = last_day_of_month.replace(day=1) - datetime.timedelta(days=1)
        remaining_days = (last_day_of_month - today).days

        forecasted_cost = daily_average * remaining_days
        
        # Get current month-to-date cost
        query_mtd = f"""
            SELECT SUM(cost) as total_cost
            FROM `{self.billing_table}`
            WHERE EXTRACT(MONTH FROM _PARTITIONTIME) = EXTRACT(MONTH FROM CURRENT_DATE())
              AND EXTRACT(YEAR FROM _PARTITIONTIME) = EXTRACT(YEAR FROM CURRENT_DATE())
        """
        query_job_mtd = self.bq_client.query(query_mtd)
        results_mtd = query_job_mtd.result()
        mtd_cost = 0
        for row in results_mtd:
            mtd_cost = row.total_cost or 0
        
        total_forecast = mtd_cost + forecasted_cost

        return {
            "daily_average_last_7_days": daily_average,
            "month_to_date_cost": mtd_cost,
            "forecasted_remaining_cost": forecasted_cost,
            "total_monthly_forecast": total_forecast,
        }

if __name__ == '__main__':
    agent = BillingAgent()
    print("Summary:", agent.get_summary())
    print("Forecast:", agent.get_forecast())