import streamlit as st
import pandas as pd
from huggingface_hub import InferenceClient
import re

# Initialize hosted inference client
client = InferenceClient(model="google/flan-t5-base")

# Simulated chart of accounts mapping
account_map = {
    "rent": "60001",
    "utilities": "60002",
    "capital": "30000",
    "cash": "10001",
    "bank": "10002",
    "sales": "40001",
    "supplies": "50001",
    "salary": "50002"
}

# Simulated business segments
segment = {
    "company": "01",
    "business_type": "102",
    "location": "001",
    "cost_center": "001",
    "future": "000"
}

# Session state to store entries
if "gl_entries" not in st.session_state:
    st.session_state.gl_entries = []

def parse_prompt(prompt):
    return client.text_generation(prompt=f"Extract accounting entry: {prompt}", max_new_tokens=50).strip()

def handle_gl_entry(prompt):
    prompt_lower = prompt.lower()
    amount = 0
    account_name = ""

    # Extract amount using regex
    amount_match = re.search(r'(\d{1,3}(,\d{3})*|\d+)(\.\d{1,2})?', prompt)
    if amount_match:
        amount = float(amount_match.group().replace(',', ''))

    # Identify transaction type
    if any(word in prompt_lower for word in ["invest", "capital", "start"]):
        account_name = "capital"
        description = "Owner Capital Contribution"
        debit_account = "cash"
        credit_account = account_name
    elif "rent" in prompt_lower:
        account_name = "rent"
        description = "Rent Expense"
        debit_account = account_name
        credit_account = "cash"
    elif "utilities" in prompt_lower:
        account_name = "utilities"
        description = "Utilities Expense"
        debit_account = account_name
        credit_account = "cash"
    elif any(word in prompt_lower for word in ["sale", "revenue"]):
        account_name = "sales"
        description = "Sales Revenue"
        debit_account = "cash"
        credit_account = account_name
    elif "supplies" in prompt_lower:
        account_name = "supplies"
        description = "Supplies Purchase"
        debit_account = account_name
        credit_account = "cash"
    elif "salary" in prompt_lower or "payroll" in prompt_lower:
        account_name = "salary"
        description = "Salary Expense"
        debit_account = account_name
        credit_account = "cash"
    else:
        description = "Unrecognized Entry"
        return pd.DataFrame([{"Date": "2025-04-01", "Description": description, "Account Code": "N/A", "Debit": 0, "Credit": 0}])

    debit_code = f"{segment['company']}-{segment['business_type']}-{segment['location']}-{segment['cost_center']}-{account_map[debit_account]}-{segment['future']}"
    credit_code = f"{segment['company']}-{segment['business_type']}-{segment['location']}-{segment['cost_center']}-{account_map[credit_account]}-{segment['future']}"

    entry = [
        {
            "Date": "2025-04-01",
            "Description": description,
            "Account Code": debit_code,
            "Debit": amount,
            "Credit": 0
        },
        {
            "Date": "2025-04-01",
            "Description": f"Offset for {description.lower()}",
            "Account Code": credit_code,
            "Debit": 0,
            "Credit": amount
        }
    ]
    st.session_state.gl_entries.extend(entry)
    return pd.DataFrame(entry)

# Streamlit UI
st.title("AI ERP Chat - MVP")
prompt = st.text_input("Enter your accounting instruction:")

delete_records = st.button("Delete All Records")
if delete_records:
    st.session_state.gl_entries = []
    st.success("All records have been deleted.")

if prompt:
    result = handle_gl_entry(prompt)
    st.dataframe(result)

# Show saved entries
if st.session_state.gl_entries:
    st.subheader("All Recorded Entries")
    st.dataframe(pd.DataFrame(st.session_state.gl_entries))