import os
import json
import pandas as pd
import time
import traceback
from openpyxl import load_workbook
from openpyxl.styles import numbers
from datetime import datetime

BASE_DIR = os.path.dirname(__file__)
CATEGORIES_FILE = os.path.join(BASE_DIR, "categories.json")

DEFAULT_CATEGORIES = {
    "Foods": [], "Beverages": [], "Cleaning": [], "Utilities": [], "Rent": [],
    "Interest Loans": [], "Advertising/Print/Decor": [], "Payroll": [],
    "Credit Fees": [], "Insurance": [], "Accounting": [], "Equipment/Furniture": [],
    "Repair/Maintenance": [], "License": [], "Misc": [], "Donations/Charity": [],
    "IRS": []
}

def load_or_initialize_categories():
    if not os.path.exists(CATEGORIES_FILE):
        with open(CATEGORIES_FILE, 'w') as file:
            json.dump(DEFAULT_CATEGORIES, file, indent=4)
        return DEFAULT_CATEGORIES
    with open(CATEGORIES_FILE, 'r') as file:
        return json.load(file)

def save_categories(categories):
    with open(CATEGORIES_FILE, 'w') as file:
        json.dump(categories, file, indent=4)

def get_user_input(categories, recorder_name):
    date = input("Enter the date (MM/DD/YYYY): ")
    datetime_obj = datetime.strptime(date, "%m/%d/%Y")
    year = datetime_obj.strftime("%Y")
    month = datetime_obj.strftime("%B")

    print("Categories:")
    for idx, cat in enumerate(categories.keys(), 1):
        print(f"{idx}. {cat}")
    category_choice = int(input("Select a category by number: "))
    category = list(categories.keys())[category_choice - 1]

    print(f"Titles in {category}:")
    if categories[category]:
        for idx, title in enumerate(categories[category], 1):
            print(f"{idx}. {title}")
        title_choice = input("Select a title by number or enter a new one: ")
        if title_choice.isdigit():
            title = categories[category][int(title_choice) - 1]
        else:
            title = title_choice
            if title not in categories[category]:
                categories[category].append(title)
    else:
        title = input("Enter the first title for this category: ")
        categories[category].append(title)

    amount = float(input("Enter the amount: "))
    addl_notes = input("Enter additional notes (if any): ")

    return {"Date": date, "Recorder": recorder_name, "Category": category, "Title": title, "Amount": amount, "Addl. Notes": addl_notes}, year, month

def set_summary_formulas(ws):
    # Define headers for the summary columns
    ws['G1'] = "Total All"
    ws['H1'] = "COGS Amount"
    ws['I1'] = "COGS %"
    ws['J1'] = "OPEX Amount"
    ws['K1'] = "OPEX %"
    ws['L1'] = "Labor Amount"
    ws['M1'] = "Labor %"

    # Total All formula
    ws['G2'] = f"=SUM(E2:E{ws.max_row})"
    ws['G2'].number_format = numbers.FORMAT_CURRENCY_USD_SIMPLE

    # COGS related formulas
    ws['H2'] = f"=SUMIF(C2:C{ws.max_row}, \"Foods\", E2:E{ws.max_row}) + SUMIF(C2:C{ws.max_row}, \"Beverages\", E2:E{ws.max_row})"
    ws['I2'] = f"=H2/G2"

    # OPEX related formulas
    opex_categories = ["Cleaning", "Utilities", "Rent", "Interest Loans", "Advertising",
                       "Credit Fees", "Insurance", "Accounting", "Equipment", "Repair", "License", "Misc", "Donations"]
    opex_formula = " + ".join([f'SUMIF(C2:C{ws.max_row}, "{cat}", E2:E{ws.max_row})' for cat in opex_categories])
    ws['J2'] = f"=({opex_formula})"
    ws['K2'] = f"=J2/G2"

    # Labor related formulas
    ws['L2'] = f"=SUMIF(C2:C{ws.max_row}, \"Payroll\", E2:E{ws.max_row}) + SUMIF(C2:C{ws.max_row}, \"IRS\", E2:E{ws.max_row})"
    ws['M2'] = f"=L2/G2"

    # Apply number formatting for financial and percentage columns
    for col in ['H2', 'J2', 'L2']:
        ws[col].number_format = numbers.FORMAT_CURRENCY_USD_SIMPLE
    for col in ['I2', 'K2', 'M2']:
        ws[col].number_format = numbers.FORMAT_PERCENTAGE_00


def ensure_directories(year, month, category, title):
    paths = {
        "month_dir": os.path.join(BASE_DIR, year, month),
        "category_dir_month": os.path.join(BASE_DIR, year, month, "Categories", category),
        "title_dir_month": os.path.join(BASE_DIR, year, month, "Categories", category, "Titles", title),
        "year_dir": os.path.join(BASE_DIR, year),
        "category_dir_year": os.path.join(BASE_DIR, year, "Categories", category),
        "title_dir_year": os.path.join(BASE_DIR, year, "Categories", category, "Titles", title)
    }
    for path in paths.values():
        os.makedirs(path, exist_ok=True)
    return paths

def update_excel(file_path, data, is_overall_summary):
    file_path = Path(file_path)
    os.makedirs(file_path.parent, exist_ok=True)

    mode = 'a' if file_path.exists() else 'w'
    sheet_name = 'Sheet1'

    with pd.ExcelWriter(file_path, engine='openpyxl', mode=mode, if_sheet_exists='overlay') as writer:
        if mode == 'a':
            book = load_workbook(file_path)
            if sheet_name in book.sheetnames:
                start_row = book[sheet_name].max_row
            else:
                start_row = 0
        else:
            start_row = 0

        df = pd.DataFrame([data])
        df.to_excel(writer, sheet_name=sheet_name, index=False, header=(start_row == 0), startrow=start_row)

    if is_overall_summary:
        wb = load_workbook(file_path)
        ws = wb[sheet_name]
        set_summary_formulas(ws)
        wb.save(file_path)

def apply_column_formatting(ws):
    # Set column widths
    for col_letter in 'ABCDEFGHIJKLMN':
        ws.column_dimensions[col_letter].width = 22

    # Format rows starting from the second
    for row in ws.iter_rows(min_row=2, max_row=ws.max_row):
        row[0].number_format = 'MM/DD/YYYY'  # Date format
        row[4].number_format = numbers.FORMAT_CURRENCY_USD_SIMPLE  # Currency format for the Amount column
        # Extend this section if more specific formatting is needed for other columns




def main():
    categories = load_or_initialize_categories()
    recorder_name = input("Enter the recorder's name: ")
    continue_recording = 'yes'

    while continue_recording.lower() == 'yes':
        try:
            data, year, month = get_user_input(categories, recorder_name)
            save_categories(categories)

            paths = ensure_directories(year, month, data["Category"], data["Title"])

            # File paths
            monthly_summary_file = os.path.join(paths["month_dir"], f'{month}_Monthly_Summary.xlsx')
            category_summary_file_month = os.path.join(paths["category_dir_month"], f'{data["Category"]}_Category_Summary.xlsx')
            title_summary_file_month = os.path.join(paths["title_dir_month"], f'{data["Title"]}_Title_Summary.xlsx')

            yearly_summary_file = os.path.join(paths["year_dir"], f'{year}_Yearly_Summary.xlsx')
            category_summary_file_year = os.path.join(paths["category_dir_year"], f'{data["Category"]}_Year_Category_Summary.xlsx')
            title_summary_file_year = os.path.join(paths["title_dir_year"], f'{data["Title"]}_Year_Title_Summary.xlsx')

            # Update Excel files with a delay to avoid conflicts
            files_to_update = [
                (monthly_summary_file, True),
                (category_summary_file_month, False),
                (title_summary_file_month, False),
                (yearly_summary_file, True),
                (category_summary_file_year, False),
                (title_summary_file_year, False)
            ]

            for file_path, is_overall in files_to_update:
                update_excel(file_path, data, is_overall_summary=is_overall)

        except Exception as e:
            print("An error occurred during the update process:")
            print(e)
            traceback.print_exc()  # To print the stack trace and understand where the error occurred

        continue_recording = input("Would you like to record another expense? (yes/no): ")

if __name__ == "__main__":
    main()