import streamlit as st
from datetime import datetime
from fpdf import FPDF # For PDF generation

# Basic page configuration
st.set_page_config(page_title="Construction Invoice Generator", layout="centered")

def create_invoice_pdf(details):
    pdf = FPDF()
    pdf.add_page()
    pdf.set_font("Arial", "B", 16)

    # --- Company Details (Top Left) & Invoice Title (Top Right) ---
    # Company Name
    pdf.set_font("Arial", "B", 14) # Slightly smaller for company name if it's long
    pdf.cell(100, 8, details['company_name'], ln=False, align="L")
    # Invoice Title
    pdf.set_font("Arial", "B", 18)
    pdf.cell(90, 8, "INVOICE", ln=True, align="R")

    # Company Address
    pdf.set_font("Arial", "", 9)
    # Use multi_cell for address in case it's long
    company_address_lines = details['company_address'].split('\n')
    company_x = pdf.get_x()
    company_y = pdf.get_y()

    for i, line in enumerate(company_address_lines):
        pdf.set_xy(company_x, company_y + (i * 4))
        pdf.cell(100, 4, line, ln=False, align="L")

    # Invoice Number (aligned with company address lines)
    pdf.set_font("Arial", "", 10)
    pdf.set_xy(company_x + 100, company_y) # Reset X to align with Invoice title column
    pdf.cell(90, 5, f"Invoice #: {details['invoice_number']}", ln=True, align="R")

    # Company Phone (below address)
    pdf.set_xy(company_x, company_y + (len(company_address_lines) * 4))
    pdf.cell(100, 5, f"Phone: {details['company_phone']}", ln=False, align="L")
    # Invoice Date (aligned with Invoice #)
    pdf.set_xy(company_x + 100, company_y + 5)
    pdf.cell(90, 5, f"Date: {details['invoice_date']}", ln=True, align="R")

    # Company Email (below phone)
    pdf.set_xy(company_x, company_y + (len(company_address_lines) * 4) + 5)
    pdf.cell(100, 5, f"Email: {details['company_email']}", ln=False, align="L")
    # Due Date (aligned with Invoice Date)
    if details.get('due_date'):
        pdf.set_xy(company_x + 100, company_y + 10)
        pdf.cell(90, 5, f"Due Date: {details['due_date']}", ln=True, align="R")
    else: # If no due date, ensure we move down appropriately
        pdf.set_xy(company_x + 100, company_y + 10) # Keep alignment
        pdf.ln(5) # Move to next line

    pdf.ln(10) # Add a break

    # --- Client Details ---
    pdf.set_font("Arial", "B", 11)
    pdf.cell(0, 7, "Bill To:", ln=True)
    pdf.set_font("Arial", "", 10)
    pdf.cell(0, 5, details['client_name'], ln=True)
    # Use multi_cell for client address
    client_address_lines = details['client_address'].split('\n')
    for line in client_address_lines:
        pdf.cell(0, 5, line, ln=True)

    if details.get('client_phone'):
        pdf.cell(0, 5, f"Phone: {details['client_phone']}", ln=True)
    if details.get('client_email'):
        pdf.cell(0, 5, f"Email: {details['client_email']}", ln=True)

    pdf.ln(5)

    # --- Project Details ---
    if details.get('project_name') or details.get('site_address'):
        pdf.set_font("Arial", "B", 11)
        pdf.cell(0, 7, "Project Details:", ln=True)
        pdf.set_font("Arial", "", 10)
        if details.get('project_name'):
            pdf.cell(0, 5, f"Project: {details['project_name']}", ln=True)
        if details.get('site_address'):
            site_address_lines = details['site_address'].split('\n')
            for line in site_address_lines:
                 pdf.cell(0, 5, f"Site: {line}", ln=True)
        pdf.ln(5)

    # --- Items Table Header ---
    pdf.set_font("Arial", "B", 10)
    pdf.set_fill_color(220, 220, 220) # Light grey background
    header_height = 7
    col_widths = {"desc": 95, "qty": 20, "unit_price": 35, "total": 40} # Adjusted for A4 total 190

    pdf.cell(col_widths["desc"], header_height, "Description", border=1, fill=True, align="C")
    pdf.cell(col_widths["qty"], header_height, "Qty", border=1, fill=True, align="C")
    pdf.cell(col_widths["unit_price"], header_height, "Unit Price", border=1, fill=True, align="C")
    pdf.cell(col_widths["total"], header_height, "Total", border=1, fill=True, ln=True, align="C")

    # --- Items Table Rows ---
    pdf.set_font("Arial", "", 9) # Slightly smaller font for item details
    subtotal = 0
    for item in details['items']:
        item_total = item['quantity'] * item['unit_price']
        subtotal += item_total

        # Calculate number of lines needed for description and determine row height
        # This is a simplified approach; for perfect alignment, more complex calculations might be needed
        # if descriptions are very long or contain many newlines.
        desc_lines = pdf.multi_cell(col_widths["desc"], 5, item['description'], border=0, align='L', dry_run=True, output='N')
        row_height = max(5 * desc_lines, 5) # Minimum height of one line (5mm)

        # Save current Y position to draw all cells of the row at the same height
        current_y = pdf.get_y()
        current_x = pdf.get_x()

        # Description (MultiCell)
        pdf.multi_cell(col_widths["desc"], 5, item['description'], border='LR', align='L')
        # After MultiCell, Y position is at the end of the multicell. X is at the beginning of the next line.
        # We need to set X, Y for the next cells in the same row.
        pdf.set_xy(current_x + col_widths["desc"], current_y)

        pdf.cell(col_widths["qty"], row_height, str(item['quantity']), border='R', align="R") # Border R to connect with Desc's R
        pdf.cell(col_widths["unit_price"], row_height, f"{item['unit_price']:.2f}", border='R', align="R")
        pdf.cell(col_widths["total"], row_height, f"{item_total:.2f}", border='R', ln=True, align="R")

    # --- Draw bottom line for the table ---
    pdf.cell(sum(col_widths.values()), 0, "", border="T", ln=True)
    pdf.ln(2)

    # --- Totals Section ---
    summary_col_width1 = col_widths["desc"] + col_widths["qty"] + col_widths["unit_price"] - 5 # width for labels
    summary_col_width2 = col_widths["total"] + 5 # width for amounts

    pdf.set_font("Arial", "", 10)
    pdf.cell(summary_col_width1, 7, "Subtotal:", align="R") # Adjusted height
    pdf.cell(summary_col_width2, 7, f"{subtotal:.2f}", border=1, ln=True, align="R")

    tax_amount = 0
    if details['tax_rate'] > 0:
        tax_amount = (subtotal * details['tax_rate']) / 100
        pdf.cell(summary_col_width1, 7, f"Tax ({details['tax_rate']}%):", align="R")
        pdf.cell(summary_col_width2, 7, f"{tax_amount:.2f}", border=1, ln=True, align="R")

    grand_total = subtotal + tax_amount
    pdf.set_font("Arial", "B", 12)
    pdf.cell(summary_col_width1, 8, "Grand Total:", align="R") # Adjusted height
    pdf.set_font("Arial", "B", 12)
    pdf.cell(summary_col_width2, 8, f"{grand_total:.2f}", border=1, ln=True, align="R")

    pdf.ln(10)

    # --- Payment Terms & Notes ---
    if details.get('payment_terms'):
        pdf.set_font("Arial", "B", 10)
        pdf.cell(0, 6, "Payment Terms:", ln=True) # Adjusted height
        pdf.set_font("Arial", "", 9)
        pdf.multi_cell(0, 5, details['payment_terms'], ln=True) # Adjusted height
        pdf.ln(3)

    if details.get('notes'):
        pdf.set_font("Arial", "B", 10)
        pdf.cell(0, 6, "Notes:", ln=True) # Adjusted height
        pdf.set_font("Arial", "", 9)
        pdf.multi_cell(0, 5, details['notes'], ln=True) # Adjusted height

    # Output the PDF: 'S' means string output, encode to latin-1 for Streamlit download
    return pdf.output(dest='S').encode('latin-1')


# --- Streamlit App UI ---
st.title("🚧 Construction Invoice Generator")

# Initialize session state for line items if it doesn't exist
if 'items' not in st.session_state:
    st.session_state.items = [] # Each item will be a dictionary

# --- Company & Client Information using columns for better layout ---
st.header("Company & Client Information")
col1, col2 = st.columns(2)

with col1:
    st.subheader("Your Company Details")
    company_name = st.text_input("Company Name", "Your Construction Co.")
    company_address = st.text_area("Company Address", "123 Main St\nAnytown, USA 12345") # Example with newline
    company_phone = st.text_input("Company Phone", "+1-555-1234")
    company_email = st.text_input("Company Email", "contact@yourconstruction.com")

with col2:
    st.subheader("Client Details")
    client_name = st.text_input("Client Name", "John Doe")
    client_address = st.text_area("Client Address", "456 Project Ave\nSometown, USA 67890") # Example with newline
    client_phone = st.text_input("Client Phone (Optional)", "")
    client_email = st.text_input("Client Email (Optional)", "")

st.markdown("---") # Visual separator

st.header("Invoice & Project Details")
col_meta1, col_meta2 = st.columns(2)
with col_meta1:
    invoice_number = st.text_input("Invoice Number", f"INV-{datetime.now().strftime('%Y%m%d%H%M')}") # Auto-generates unique-ish
    invoice_date = st.date_input("Invoice Date", datetime.today())
    due_date_option = st.checkbox("Add Due Date?", value=True)
    due_date = None
    if due_date_option:
        due_date = st.date_input("Due Date", datetime.today() + timedelta(days=30))


with col_meta2:
    project_name = st.text_input("Project Name / ID (Optional)", "Residential Build - Lot 7")
    site_address = st.text_area("Project Site Address (Optional)", "789 Site Rd\nSometown, USA 67890")

st.markdown("---")

st.header("Invoice Items")

# --- Form for adding new items ---
with st.expander("➕ Add New Item", expanded=True):
    item_cols = st.columns([3, 1, 1]) # Description, Quantity, Unit Price
    with item_cols[0]:
        new_item_description = st.text_input("Description", key="new_desc_val") # Unique keys
    with item_cols[1]:
        new_item_quantity = st.number_input("Quantity", min_value=0.01, step=0.01, format="%.2f", key="new_qty_val")
    with item_cols[2]:
        new_item_unit_price = st.number_input("Unit Price ($)", min_value=0.00, step=0.01, format="%.2f", key="new_price_val")

    if st.button("Add Item", key="add_item_button"): # Unique key for button
        if new_item_description and new_item_quantity > 0: # Unit price can be 0
            st.session_state.items.append({
                "description": new_item_description,
                "quantity": new_item_quantity,
                "unit_price": new_item_unit_price
            })
            st.success(f"Added: {new_item_description}")
            # To clear inputs, you'd typically use form submission or more complex state management.
            # For simplicity, we'll let users clear them manually or type over.
        else:
            st.error("Please fill in item description and quantity.")

# --- Display existing items and allow removal ---
if st.session_state.items:
    st.subheader("Current Items:")
    subtotal = 0
    for i, item in enumerate(st.session_state.items):
        item_total = item['quantity'] * item['unit_price']
        subtotal += item_total
        cols = st.columns([0.5, 0.15, 0.15, 0.15, 0.05]) # Desc, Qty, Unit Price, Total, Remove
                                                    # Weights for column widths
        cols[0].write(item['description'])
        cols[1].write(f"{item['quantity']:.2f}")
        cols[2].write(f"${item['unit_price']:.2f}")
        cols[3].write(f"${item_total:.2f}")
        if cols[4].button(f"🗑️", key=f"remove_item_{i}"): # Using an icon for remove
            st.session_state.items.pop(i)
            st.rerun()
    st.markdown(f"**Subtotal: ${subtotal:.2f}**")
else:
    st.info("No items added yet.")

st.markdown("---")

st.header("Summary & Payment")
tax_rate = st.number_input("Tax Rate (%)", min_value=0.0, max_value=100.0, value=0.0, step=0.1, format="%.1f") # Default 0% tax

# Recalculate totals based on current items for display
current_subtotal = sum(item['quantity'] * item['unit_price'] for item in st.session_state.items)
tax_amount = (current_subtotal * tax_rate) / 100
grand_total = current_subtotal + tax_amount

st.markdown(f"#### Calculated Subtotal: ${current_subtotal:,.2f}")
st.markdown(f"#### Tax Amount: ${tax_amount:,.2f}")
st.markdown(f"### Grand Total: ${grand_total:,.2f}")


st.markdown("---")
payment_terms = st.text_area("Payment Terms", "Payment due within 30 days.\nBank Transfer: Account #12345, Sort Code: 00-00-00, Bank of Construction.")
notes = st.text_area("Additional Notes/Comments (Optional)", "Thank you for your business! We appreciate your prompt payment.")

st.markdown("---")

if st.button("🚀 Generate Invoice PDF", type="primary", use_container_width=True):
    if not st.session_state.items:
        st.error("⚠️ Please add at least one item to the invoice.")
    elif not company_name or not client_name or not invoice_number:
        st.error("⚠️ Please fill in all required company, client, and invoice number details.")
    else:
        invoice_details = {
            "company_name": company_name,
            "company_address": company_address,
            "company_phone": company_phone,
            "company_email": company_email,
            "client_name": client_name,
            "client_address": client_address,
            "client_phone": client_phone,
            "client_email": client_email,
            "invoice_number": invoice_number,
            "invoice_date": invoice_date.strftime("%B %d, %Y"), # Nicer date format
            "due_date": due_date.strftime("%B %d, %Y") if due_date else "",
            "project_name": project_name,
            "site_address": site_address,
            "items": st.session_state.items,
            "tax_rate": tax_rate,
            "payment_terms": payment_terms,
            "notes": notes
        }

        try:
            with st.spinner("Generating PDF... Please wait."): # Show a spinner
                pdf_data = create_invoice_pdf(invoice_details)
            st.success("🎉 Invoice PDF Generated Successfully!")

            pdf_filename = f"Invoice_{invoice_details['invoice_number'].replace('/', '-')}_{invoice_details['client_name'].replace(' ', '_')}.pdf"

            st.download_button(
                label="📥 Download Invoice PDF",
                data=pdf_data,
                file_name=pdf_filename,
                mime="application/pdf",
                use_container_width=True
            )
        except Exception as e:
            st.error(f"An error occurred during PDF generation: {e}")
            st.error("Details: Please ensure all text inputs (like addresses, descriptions) are reasonably formatted for PDF generation. Long unbroken strings or unusual characters can sometimes cause issues with FPDF.")

# Add a little footer
st.markdown("---")
st.markdown("<div style='text-align: center; color: grey;'>Construction Invoice Generator by Your AI Teacher</div>", unsafe_allow_html=True)