import os
import boto3
import datetime
import logging

class AWSBillingAgent:
    def __init__(self):
        self.aws_access_key_id = os.environ.get("AWS_ACCESS_KEY_ID")
        self.aws_secret_access_key = os.environ.get("AWS_SECRET_ACCESS_KEY")
        self.aws_default_region = os.environ.get("AWS_DEFAULT_REGION", "us-east-1")

        if not all([self.aws_access_key_id, self.aws_secret_access_key, self.aws_default_region]):
            raise ValueError("AWS credentials (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_DEFAULT_REGION) must be set.")

        self.ce_client = boto3.client(
            'ce',
            aws_access_key_id=self.aws_access_key_id,
            aws_secret_access_key=self.aws_secret_access_key,
            region_name=self.aws_default_region
        )
        logging.basicConfig(level=logging.INFO)

    def get_summary(self):
        """
        Henter MTD-kostnad, daglig snitt, prognose og en fordeling per tjeneste for AWS.
        """
        try:
            today = datetime.date.today()
            start_of_month = today.replace(day=1).isoformat()
            end_of_month = today.isoformat()

            # Hent kostnad per tjeneste
            response = self.ce_client.get_cost_and_usage(
                TimePeriod={
                    'Start': start_of_month,
                    'End': end_of_month
                },
                Granularity='MONTHLY',
                Metrics=['UnblendedCost'],
                GroupBy=[{'Type': 'DIMENSION', 'Key': 'SERVICE'}]
            )

            summary = []
            for item in response['ResultsByTime'][0]['Groups']:
                summary.append({
                    "service": item['Keys'][0],
                    "total_cost": float(item['Metrics']['UnblendedCost']['Amount'])
                })

            summary = sorted(summary, key=lambda i: i['total_cost'], reverse=True)
            
            if not summary:
                return {
                    "onboarding_status": {
                        "state": "awaiting_data",
                        "message": "Kunne ikke hente faktureringsdata fra AWS. Sjekk at Cost Explorer er aktivert."
                    }
                }
            
            return {"summary": summary}
        except Exception as e:
            logging.error(f"Error fetching AWS summary: {e}")
            return {"error": str(e)}

    def get_forecast(self):
        """
        Henter en kostnadsprognose for AWS.
        """
        try:
            today = datetime.date.today()
            start_of_month = today.replace(day=1).isoformat()
            
            # MTD-kostnad
            mtd_response = self.ce_client.get_cost_and_usage(
                TimePeriod={
                    'Start': start_of_month,
                    'End': today.isoformat()
                },
                Granularity='MONTHLY',
                Metrics=['UnblendedCost']
            )
            mtd_cost = float(mtd_response['ResultsByTime'][0]['Total']['UnblendedCost']['Amount'])

            # Prognose
            forecast_response = self.ce_client.get_cost_forecast(
                TimePeriod={
                    'Start': (today + datetime.timedelta(days=1)).isoformat(),
                    'End': (today.replace(day=1) + datetime.timedelta(days=32)).replace(day=1).isoformat()
                },
                Metric='UNBLENDED_COST',
                Granularity='MONTHLY'
            )
            total_forecast = float(forecast_response['Total']['Amount'])

            # Daglig snitt
            last_7_days_start = (today - datetime.timedelta(days=7)).isoformat()
            daily_avg_response = self.ce_client.get_cost_and_usage(
                TimePeriod={
                    'Start': last_7_days_start,
                    'End': today.isoformat()
                },
                Granularity='DAILY',
                Metrics=['UnblendedCost']
            )
            
            daily_costs = [float(item['Total']['UnblendedCost']['Amount']) for item in daily_avg_response['ResultsByTime']]
            daily_average = sum(daily_costs) / len(daily_costs) if daily_costs else 0

            return {
                "daily_average_last_7_days": daily_average,
                "month_to_date_cost": mtd_cost,
                "total_monthly_forecast": total_forecast,
                "data_note": "Prognose fra AWS Cost Explorer. Kan avvike fra endelig faktura."
            }

        except Exception as e:
            logging.error(f"Error fetching AWS forecast: {e}")
            return {"error": str(e)}

if __name__ == '__main__':
    # For local testing, ensure env vars are set
    agent = AWSBillingAgent()
    print("Summary:", agent.get_summary())
    print("Forecast:", agent.get_forecast())