from datetime import datetime
from app.utils.calc_api_usage import calculate_api_usage
from flask import current_app
from typing import List, Dict, Any, Optional 
from bson import ObjectId

class MongoBillingService:
    def __init__(self):
        self.mongo_client = current_app.mongo_client 

    def update_usage_total_cost(self,  
                            product: str, 
                            start_date: str, 
                            end_date: str,
                            price: float,
                            billing_unit: int,
                            company_ids: Optional[List[str]] = None, ) -> int:
        query = { 
            "product": product,
            "createdAt": {
                "$gte": datetime.strptime(f"{start_date} 00:00:00", "%Y-%m-%d %H:%M:%S"),
                "$lte": datetime.strptime(f"{end_date} 23:59:59", "%Y-%m-%d %H:%M:%S")
            }
        } 

        if company_ids:
            query["companyId"] = {"$in": company_ids}

        collection = self.mongo_client["billing-api"]["api_usages"]

        cursor = collection.find(query) 

        count = 0 

        for doc in cursor:
            usage = float(doc.get("usage", 0))
            new_total_cost = calculate_api_usage(float(price), int(billing_unit), float(usage))

            collection.update_one(
                {"_id": doc["_id"]},
                {
                    "$set": {
                        "total_cost": new_total_cost,
                        "price": f"{price}",
                        "billingUnit": billing_unit,
                        "updatedAt": datetime.utcnow()
                    }
                }
            )
            count+=1
        return count

    def update_model_policy_price(self, id:str, update_data: Dict[str, Any]) -> int:

        collection = self.mongo_client["billing-api"]["api_pricings"]

        update_data["updatedAt"] = datetime.utcnow()

        result = collection.update_one(
            {"_id": ObjectId(id)},
            {"$set": update_data}
        )
  
        return result.modified_count