from pydantic import BaseModel
from typing import Annotated
from .clock import Clock
from autogen import AssistantAgent
import pandas as pd
from .app import App
from .object_store import ObjectStore

class BudgetApp(App):
    funds: float = 0
    accounts: dict = {}
    history: list[dict] = []
    clock: Clock
    store: ObjectStore
    
    def __init__(self, clock: Clock, store: ObjectStore):
        self.clock = clock
        self.store = store
    
    def set_available_funds(self, amount: float) -> str:
        self.accounts["Not Allocated"] = amount
        self.on_change()
        return "Allocated %.2f to Not Allocated" % amount
    
    def allocate(self, 
                 category: Annotated[str, "The category to which allocate funds" ], 
                 amount: Annotated[float, "The amount to allocate"]) -> str:
        
        self.accounts["Not Allocated"] > amount
        self.accounts["Not Allocated"] -= amount
        if self.accounts.get(category):
            self.accounts[category] += amount
        else:
            self.accounts[category] = amount
        self.on_change()
        return "Allocated"

    def get_avaliable_funds(self, 
                category: Annotated[str, "the category to get the available funds. If not provided it will return all funds available" ] = None) -> float:
        if category is None:
            return self.get_total_avaliable_funds()
        return self.accounts[category]
    
    def get_allocations(self) -> dict:
        return self.accounts
    
    def get_current_budget(self) -> str:
        if len(self.accounts) == 0:
            return ""
        allocations = "Current Budget:"
        for category, amount in self.accounts.items():
            allocations += f"{category}: {amount}\n"
        return allocations
        
    def get_total_avaliable_funds(self) -> float:
        return sum(self.accounts.values())
    
    
    
    def spend(self, amount: float,  category: str = None) -> str:
        provisioned = self.accounts.get(category)
        if category not in self.accounts:
            self.accounts["Not Allocated"] -= amount
        else:
            self.accounts[category] -= amount
        
        self.history.append({
            "date": self.clock.now().isoformat(),
            "category": category,
            "amount": amount * -1,
            "type": "spend"
        })
        self.on_change()
        if provisioned is None:
            return "Spent %.2f from Not Allocated as %s was not provisioned" % (amount,category)
        return "Spent %.2f from %s" % (amount, category)
        
    def deposit(self, category: str, amount: float) -> str:
        if category not in self.accounts:
            self.accounts[category] = 0
        self.accounts[category] += amount    
        self.history.append({
            "date": self.clock.now().isoformat(),
            "category": category,
            "amount": amount,
            "type": "deposit"
        })
        self.on_change()
        return "Deposited %.2f to %s" % (amount, category)
    
    def on_change(self):
        self.store.save("budget", {
            "funds": self.funds,
            "accounts": self.accounts,
            "history": self.history
        }, namespace="budget")
    
    def get_cashflow(self, 
                     start_date: Annotated[str, "The start date"], 
                     end_date: Annotated[str, "The end date"] ) -> str:
        df = pd.DataFrame(self.history)
        df = df[(df["date"] >= start_date) & (df["date"] <= end_date)]
        return df.groupby("category").sum("amount").to_markdown()
    
    def serialize(self) -> str:
        return self.model_dump_json()
    
    def get_functions(self):
        def set_available_funds(amount: float) -> str:
            return self.set_available_funds(amount)
        
        def allocate(category: str, amount: float) -> str:
            return self.allocate(category, amount)
        
        def get_avaliable_funds(category: str) -> float:
            return self.get_avaliable_funds(category)
        
        def get_total_avaliable_funds() -> float:
            return self.get_total_avaliable_funds()
        
        def spend(category: str, amount: float) -> str:
            return self.spend(amount, category)
        
        def deposit(category: str, amount: float) -> str:
            return self.deposit(category, amount)
        
        def get_cashflow(start_date: str, end_date: str) -> str:
            return self.get_cashflow(start_date, end_date)
        
        def get_allocations() -> dict:
            return self.get_allocations()
        
        return [
            (set_available_funds, "Use this tool to set the available funds when starting an account and there were no funds. The amount will be allocated to the `Not Allocated` category"),
            (allocate, "Use this tool to allocate funds to a category"),
            (get_avaliable_funds, "Use this tool to get the avaliable funds for a category"),
            (get_total_avaliable_funds, "Use this tool to get the total avaliable funds"),
            (spend, "Use this tool to spend funds from a category"),
            (deposit, "Use this tool to deposit funds to a category"),
            (get_cashflow, "Use this tool to get the cashflow for a period of time"),
            (get_allocations, "Use this tool to get the allocations for the budget")
        ]
        