import sqlite3
import json
from transformers import AutoModelForCausalLM, AutoTokenizer
import uuid
import datetime

# Initialize AI model (Mistral-7B)
model_name = "mistralai/Mixtral-8x7B-Instruct-v0.1"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Database setup
conn = sqlite3.connect("erp.db")
cursor = conn.cursor()

# Create tables
cursor.execute("""
CREATE TABLE IF NOT EXISTS chart_of_accounts (
    account_id TEXT PRIMARY KEY,
    account_name TEXT,
    account_type TEXT,
    parent_id TEXT,
    allow_budgeting BOOLEAN,
    allow_posting BOOLEAN
)
""")
cursor.execute("""
CREATE TABLE IF NOT EXISTS journal_entries (
    entry_id TEXT PRIMARY KEY,
    date TEXT,
    debit_account_id TEXT,
    credit_account_id TEXT,
    amount REAL,
    description TEXT
)
""")
conn.commit()

# Debit/Credit rules
ACCOUNT_RULES = {
    "Asset": {"increase": "Debit", "decrease": "Credit"},
    "Liability": {"increase": "Credit", "decrease": "Debit"},
    "Equity": {"increase": "Credit", "decrease": "Debit"},
    "Revenue": {"increase": "Credit", "decrease": "Debit"},
    "Expense": {"increase": "Debit", "decrease": "Credit"}
}

# Sample chart of accounts
def initialize_chart_of_accounts():
    accounts = [
        ("1", "Fixed Assets", "Asset", None, True, False),
        ("1.1", "Plant", "Asset", "1", True, True),
        ("1.2", "Machinery", "Asset", "1", True, True),
        ("2", "Cash", "Asset", None, True, True),
        ("3", "Accounts Payable", "Liability", None, True, True)
    ]
    cursor.executemany("""
    INSERT OR REPLACE INTO chart_of_accounts
    (account_id, account_name, account_type, parent_id, allow_budgeting, allow_posting)
    VALUES (?, ?, ?, ?, ?, ?)
    """, accounts)
    conn.commit()

# Parse prompt using AI model
def parse_prompt(prompt):
    input_text = f"Parse the following accounting prompt and return a JSON object with debit and credit accounts, amount, and payment method (cash or credit). Prompt: {prompt}"
    inputs = tokenizer(input_text, return_tensors="pt")
    outputs = model.generate(**inputs, max_length=200)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # Assume model returns JSON-like string
    try:
        parsed = json.loads(response)
        return parsed
    except:
        # Fallback parsing
        if "laptop" in prompt.lower() and "cash" in prompt.lower():
            return {
                "debit": {"account": "Laptop", "type": "Asset", "amount": 200},
                "credit": {"account": "Cash", "type": "Asset", "amount": 200},
                "payment_method": "cash"
            }
        return None

# Generate journal entry
def generate_journal_entry(prompt, follow_up_response=None):
    parsed = parse_prompt(prompt)
    if not parsed:
        return "Unable to parse prompt. Please clarify."

    debit_account = parsed["debit"]["account"]
    credit_account = parsed["credit"]["account"]
    amount = parsed["debit"]["amount"]
    payment_method = parsed.get("payment_method")

    # Follow-up for payment method
    if not payment_method:
        return "Wonderful, did you buy on credit? (Yes/No)"

    # Adjust credit account for credit purchases
    if follow_up_response and follow_up_response.lower() == "yes":
        credit_account = "Accounts Payable"
        credit_type = "Liability"
    else:
        credit_type = parsed["credit"]["type"]

    # Find account IDs
    cursor.execute("SELECT account_id, allow_posting FROM chart_of_accounts WHERE account_name = ?", (debit_account,))
    debit_result = cursor.fetchone()
    cursor.execute("SELECT account_id, allow_posting FROM chart_of_accounts WHERE account_name = ?", (credit_account,))
    credit_result = cursor.fetchone()

    if not debit_result or not credit_result:
        return "Invalid account names."
    if not debit_result[1] or not credit_result[1]:
        return "Posting not allowed for one or both accounts."

    # Save journal entry
    entry_id = str(uuid.uuid4())
    date = datetime.datetime.now().isoformat()
    cursor.execute("""
    INSERT INTO journal_entries (entry_id, date, debit_account_id, credit_account_id, amount, description)
    VALUES (?, ?, ?, ?, ?, ?)
    """, (entry_id, date, debit_result[0], credit_result[0], amount, prompt))
    conn.commit()

    return f"Journal Entry Created: Debit {debit_account} ${amount}, Credit {credit_account} ${amount}"

# Generate T-account
def generate_t_account(account_name):
    cursor.execute("SELECT account_id FROM chart_of_accounts WHERE account_name = ?", (account_name,))
    account_id = cursor.fetchone()
    if not account_id:
        return "Account not found."
    
    account_id = account_id[0]
    cursor.execute("""
    SELECT amount, description, 'Debit' as type FROM journal_entries WHERE debit_account_id = ?
    UNION
    SELECT amount, description, 'Credit' as type FROM journal_entries WHERE credit_account_id = ?
    """, (account_id, account_id))
    entries = cursor.fetchall()
    
    t_account = f"T-Account for {account_name}\n{'-'*40}\nDebit | Credit\n{'-'*40}\n"
    for amount, desc, entry_type in entries:
        if entry_type == "Debit":
            t_account += f"${amount} |       | {desc}\n"
        else:
            t_account += f"       | ${amount} | {desc}\n"
    
    return t_account

# Example usage
if __name__ == "__main__":
    initialize_chart_of_accounts()
    
    # Test prompt
    prompt = "Bought a laptop for $200 on cash"
    result = generate_journal_entry(prompt)
    print(result)
    
    # Test T-account
    t_account = generate_t_account("Laptop")
    print(t_account)
    
    # Test follow-up
    prompt = "Bought a laptop for $200"
    result = generate_journal_entry(prompt)
    print(result)
    if "did you buy on credit" in result.lower():
        result = generate_journal_entry(prompt, "Yes")
        print(result)