from flask import Flask, render_template, request, redirect, url_for, send_from_directory, abort
import os
import mysql.connector
from werkzeug.utils import secure_filename
from config import db_config


app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = os.path.join('static', 'uploads')
ALLOWED_EXTENSIONS = {'pdf', 'docx', 'doc', 'xlsx', 'xls'}

@app.route('/')
def welcome():
    return render_template('welcome.html')

@app.route('/dashboard')
def index():
    return render_template('index.html')  # Your dashboard page

# Ensure folder exists
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)


def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS


# Upload route
@app.route('/upload', methods=['GET', 'POST'])
def upload_file():
    if request.method == 'POST':
        files = request.files.getlist('documents')
        year = request.form['year']
        stage = request.form['stage']

        conn = mysql.connector.connect(**db_config)
        cursor = conn.cursor()

        for file in files:
            if file and allowed_file(file.filename):
                filename = secure_filename(file.filename)
                filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
                file.save(filepath)

                cursor.execute("""
                    INSERT INTO documents (filename, filepath, filetype, year, stage)
                    VALUES (%s, %s, %s, %s, %s)
                """, (filename, filepath, file.filename.rsplit('.', 1)[1], year, stage))

        conn.commit()
        cursor.close()
        conn.close()
        return redirect(url_for('view_documents'))

    return render_template('upload.html')


# View all documents with filters
@app.route('/documents')
def view_documents():
    year = request.args.get('year')
    stage = request.args.get('stage')

    conn = mysql.connector.connect(**db_config)
    cursor = conn.cursor(dictionary=True)

    query = "SELECT * FROM documents WHERE 1=1"
    params = []

    if year:
        query += " AND year = %s"
        params.append(year)
    if stage:
        query += " AND stage = %s"
        params.append(stage)

    cursor.execute(query, params)
    documents = cursor.fetchall()

    cursor.execute("SELECT DISTINCT year FROM documents ORDER BY year DESC")
    years = [row['year'] for row in cursor.fetchall()]

    cursor.close()
    conn.close()

    return render_template('view_docs.html', documents=documents, years=years)


# Serve uploaded file
from flask import send_file

@app.route('/uploads/<filename>')
def uploaded_file(filename):
    mode = request.args.get('mode', 'view')
    filepath = os.path.join(app.config['UPLOAD_FOLDER'], secure_filename(filename))

    if not os.path.exists(filepath):
        abort(404)

    file_ext = filename.rsplit('.', 1)[-1].lower()

    # --- View Mode ---
    if mode == 'view':
        if file_ext == 'pdf':
            return send_file(filepath, mimetype='application/pdf')
        elif file_ext in ['xls', 'xlsx']:
            # Excel cannot be rendered in-browser by Flask; trigger download instead
            return send_file(filepath, as_attachment=False, download_name=filename, mimetype='application/vnd.openxmlformats-officedocument.spreadsheetml.sheet')
        else:
            return abort(415)  # Unsupported type for viewing

    # --- Download Mode ---
    return send_file(filepath, as_attachment=True)


# (Keep all your other routes and imports as they are)

## 1. READ/DISPLAY all ITR records
# This page will show all records in a table with Edit and Delete buttons.
@app.route('/itr_records')
def display_itr():
    conn = get_db_connection()
    cursor = conn.cursor(dictionary=True)
    cursor.execute("SELECT * FROM itr ORDER BY year DESC, id DESC")
    records = cursor.fetchall()
    cursor.close()
    conn.close()
    return render_template('display_itr.html', records=records)


## 2. CREATE/ADD a new ITR record
# This route handles both showing the blank form and saving the new data.
@app.route('/itr/add', methods=['GET', 'POST'])
def add_itr():
    if request.method == 'POST':
        conn = get_db_connection()
        cursor = conn.cursor()

        # A list of all columns in your form and database table
        columns = [
            'year', 'gross_total_income', 'disallowance_14a', 'disallowance_37',
            'deduction_80ia_business', 'deduction_80ia_misc', 'deduction_80ia_other',
            'deduction_sec37_disallowance', 'deduction_80g', 'net_taxable_income',
            'tax_30_percent', 'tax_book_profit_18_5', 'tax_payable', 'surcharge_12',
            'edu_cess_3', 'total_tax_payable', 'mat_credit', 'interest_234c',
            'total_tax', 'advance_tax', 'tds', 'tcs', 'tax_on_assessment', 'refund'
        ]

        query = f"INSERT INTO itr ({', '.join(columns)}) VALUES ({', '.join(['%s'] * len(columns))})"
        values = [request.form.get(col, 0) for col in columns]

        cursor.execute(query, tuple(values))
        conn.commit()
        cursor.close()
        conn.close()
        # After adding, redirect to the page that shows all records
        return redirect(url_for('display_itr'))

    # If it's a GET request, just show the blank form to add a record
    return render_template('add_itr.html')


## 3. UPDATE an existing ITR record
# This route needs an ID to know which record to edit.
@app.route('/itr/update/<int:id>', methods=['GET', 'POST'])
def update_itr(id):
    conn = get_db_connection()

    if request.method == 'POST':
        cursor = conn.cursor()
        columns = [
            'year', 'gross_total_income', 'disallowance_14a', 'disallowance_37',
            'deduction_80ia_business', 'deduction_80ia_misc', 'deduction_80ia_other',
            'deduction_sec37_disallowance', 'deduction_80g', 'net_taxable_income',
            'tax_30_percent', 'tax_book_profit_18_5', 'tax_payable', 'surcharge_12',
            'edu_cess_3', 'total_tax_payable', 'mat_credit', 'interest_234c',
            'total_tax', 'advance_tax', 'tds', 'tcs', 'tax_on_assessment', 'refund'
        ]

        # Create the "SET column = %s" part of the query
        set_clause = ', '.join([f"{col} = %s" for col in columns])
        query = f"UPDATE itr SET {set_clause} WHERE id = %s"

        values = [request.form.get(col, 0) for col in columns]
        values.append(id)  # Add the ID for the WHERE clause at the end

        cursor.execute(query, tuple(values))
        conn.commit()
        cursor.close()
        conn.close()
        return redirect(url_for('display_itr'))

    # For a GET request, fetch the existing data and show it in the form
    cursor = conn.cursor(dictionary=True)
    cursor.execute("SELECT * FROM itr WHERE id = %s", (id,))
    record = cursor.fetchone()
    cursor.close()
    conn.close()
    return render_template('update_itr.html', record=record)


## 4. DELETE an ITR record
# This route also needs an ID to know which record to delete.
@app.route('/itr/delete/<int:id>', methods=['POST'])
def delete_itr(id):
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute("DELETE FROM itr WHERE id = %s", (id,))
    conn.commit()
    cursor.close()
    conn.close()
    # After deleting, redirect back to the display page
    return redirect(url_for('display_itr'))


@app.route('/itr', methods=['GET', 'POST'])
def itr_form():
    if request.method == 'POST':
        data = {key: request.form.get(key, 0) for key in request.form}
        conn = mysql.connector.connect(**db_config)
        cursor = conn.cursor()
        query = """
            INSERT INTO itr (
                year, gross_total_income, disallowance_14a, disallowance_37,
                deduction_80ia_business, deduction_80ia_misc, deduction_80ia_other,
                deduction_sec37_disallowance, deduction_80g, net_taxable_income,
                tax_30_percent, tax_book_profit_18_5, tax_payable, surcharge_12,
                edu_cess_3, total_tax_payable, mat_credit, interest_234c,
                total_tax, advance_tax, tds, tcs, tax_on_assessment, refund
            ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
        """
        values = tuple([
            int(data.get('year', 0))
        ] + [
            float(data.get(col, 0)) for col in [
                'gross_total_income', 'disallowance_14a', 'disallowance_37',
                'deduction_80ia_business', 'deduction_80ia_misc', 'deduction_80ia_other',
                'deduction_sec37_disallowance', 'deduction_80g', 'net_taxable_income',
                'tax_30_percent', 'tax_book_profit_18_5', 'tax_payable', 'surcharge_12',
                'edu_cess_3', 'total_tax_payable', 'mat_credit', 'interest_234c',
                'total_tax', 'advance_tax', 'tds', 'tcs', 'tax_on_assessment', 'refund'
            ]
        ])
        cursor.execute(query, values)
        conn.commit()
        cursor.close()
        conn.close()
        return redirect(url_for('index'))
    return render_template('itr_form.html')


#
# ADD THESE NEW FUNCTIONS TO YOUR APP.PY FILE
#

## ===============================================
## AO (Assessing Officer) Routes
## ===============================================

# DISPLAY all AO records
@app.route('/ao_records')
def display_ao():
    conn = get_db_connection()
    cursor = conn.cursor(dictionary=True)
    # Note: Querying the 'ao' table
    cursor.execute("SELECT * FROM ao ORDER BY year DESC, id DESC")
    records = cursor.fetchall()
    cursor.close()
    conn.close()
    # Note: Rendering the 'display_ao.html' template
    return render_template('display_ao.html', records=records)


# ADD a new AO record
@app.route('/ao/add', methods=['GET', 'POST'])
def add_ao():
    if request.method == 'POST':
        conn = get_db_connection()
        cursor = conn.cursor()

        # Define the columns for the 'ao' table
        columns = [
            'year', 'gross_total_income', 'disallowance_14a', 'disallowance_37',
            'deduction_80ia_business', 'deduction_sec37_disallowance', 'deduction_80g',
            'net_taxable_income', 'tax_30_percent', 'tax_book_profit_18_5',
            'surcharge_12', 'edu_cess_3', 'total_tax_payable', 'mat_credit',
            'interest_234c', 'total_tax', 'advance_tax', 'tds', 'tcs',
            'tax_on_assessment', 'refund'
        ]  # Make sure these match your 'ao' table columns!

        # Note: Inserting into the 'ao' table
        query = f"INSERT INTO ao ({', '.join(columns)}) VALUES ({', '.join(['%s'] * len(columns))})"
        values = [request.form.get(col, 0) for col in columns]

        cursor.execute(query, tuple(values))
        conn.commit()
        cursor.close()
        conn.close()
        # Note: Redirecting to the 'display_ao' function
        return redirect(url_for('display_ao'))

    # Note: Rendering the 'add_ao.html' template
    return render_template('add_ao.html')


# (You will also need to add update_ao and delete_ao functions later)


@app.route('/ao', methods=['GET', 'POST'])
def ao_form():
    if request.method == 'POST':
        data = {key: request.form.get(key, 0) for key in request.form}
        conn = mysql.connector.connect(**db_config)
        cursor = conn.cursor()

        query = """
            INSERT INTO ao (
                year, gross_total_income, disallowance_14a, disallowance_37,
                deduction_80ia_business, deduction_sec37_disallowance, deduction_80g,
                net_taxable_income, tax_30_percent, tax_book_profit_18_5,
                surcharge_12, edu_cess_3, total_tax_payable, mat_credit,
                interest_234c, total_tax, advance_tax, tds, tcs,
                tax_on_assessment, refund
            ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
        """

        values = tuple([
            data.get('year'),
            float(data.get('gross_total_income', 0)),
            float(data.get('disallowance_14a', 0)),
            float(data.get('disallowance_37', 0)),
            float(data.get('deduction_80ia_business', 0)),
            float(data.get('deduction_sec37_disallowance', 0)),
            float(data.get('deduction_80g', 0)),
            float(data.get('net_taxable_income', 0)),
            float(data.get('tax_30_percent', 0)),
            float(data.get('tax_book_profit_18_5', 0)),
            float(data.get('surcharge_12', 0)),
            float(data.get('edu_cess_3', 0)),
            float(data.get('total_tax_payable', 0)),
            float(data.get('mat_credit', 0)),
            float(data.get('interest_234c', 0)),
            float(data.get('total_tax', 0)),
            float(data.get('advance_tax', 0)),
            float(data.get('tds', 0)),
            float(data.get('tcs', 0)),
            float(data.get('tax_on_assessment', 0)),
            float(data.get('refund', 0)),
        ])

        cursor.execute(query, values)
        conn.commit()
        cursor.close()
        conn.close()
        return redirect(url_for('index'))
    
    return render_template('ao_form.html')


#
# ADD THESE NEW CIT FUNCTIONS TO YOUR APP.PY FILE
#

## =======================================================
## CIT (Commissioner of Income Tax) Routes
## =======================================================

# DISPLAY all CIT records
@app.route('/cit_records')
def display_cit():
    conn = get_db_connection()
    cursor = conn.cursor(dictionary=True)
    # Querying the 'cit' table
    cursor.execute("SELECT * FROM cit ORDER BY year DESC, id DESC")
    records = cursor.fetchall()
    cursor.close()
    conn.close()
    # Rendering the 'display_cit.html' template
    return render_template('display_cit.html', records=records)


# ADD a new CIT record
@app.route('/cit/add', methods=['GET', 'POST'])
def add_cit():
    if request.method == 'POST':
        conn = get_db_connection()
        cursor = conn.cursor()

        # IMPORTANT: These columns match your 'cit' table structure
        columns = [
            'year', 'gross_total_income', 'deduction_80ia_business',
            'deduction_sec37_disallowance', 'deduction_80g', 'net_taxable_income',
            'tax_30_percent', 'tax_book_profit_18_5', 'tax_payable', 'surcharge_12',
            'edu_cess_3', 'total_tax_payable', 'mat_credit', 'interest_234c',
            'total_tax', 'advance_tax', 'tds', 'tcs', 'tax_on_assessment', 'refund'
        ]

        # Inserting into the 'cit' table
        query = f"INSERT INTO cit ({', '.join(columns)}) VALUES ({', '.join(['%s'] * len(columns))})"
        values = [request.form.get(col, 0) for col in columns]

        cursor.execute(query, tuple(values))
        conn.commit()
        cursor.close()
        conn.close()
        # Redirecting to the 'display_cit' function
        return redirect(url_for('display_cit'))

    # Rendering the 'add_cit.html' template
    return render_template('add_cit.html')


# (You will also need to add update_cit and delete_cit functions later)


#
# ADD THESE FINAL FUNCTIONS FOR ITAT TO YOUR APP.PY FILE
#

## =======================================================
## ITAT (Income Tax Appellate Tribunal) Routes
## =======================================================

# DISPLAY all ITAT records
@app.route('/itat_records')
def display_itat():
    conn = get_db_connection()
    cursor = conn.cursor(dictionary=True)
    # Querying the 'itat' table
    cursor.execute("SELECT * FROM itat ORDER BY year DESC, id DESC")
    records = cursor.fetchall()
    cursor.close()
    conn.close()
    # Rendering the 'display_itat.html' template
    return render_template('display_itat.html', records=records)


# ADD a new ITAT record
@app.route('/itat/add', methods=['GET', 'POST'])
def add_itat():
    if request.method == 'POST':
        conn = get_db_connection()
        cursor = conn.cursor()

        # NOTE: These are the specific columns for your 'itat' table
        columns = [
            'year', 'mat_tax_credit', 'surcharge', 'cess', 'total_credit'
        ]

        # Inserting into the 'itat' table
        query = f"INSERT INTO itat ({', '.join(columns)}) VALUES ({', '.join(['%s'] * len(columns))})"
        values = [request.form.get(col, 0) for col in columns]

        cursor.execute(query, tuple(values))
        conn.commit()
        cursor.close()
        conn.close()
        # Redirecting to the 'display_itat' function
        return redirect(url_for('display_itat'))

    # Rendering the 'add_itat.html' template
    return render_template('add_itat.html')


# (You will also need to add update_itat and delete_itat functions later)

@app.route('/cit', methods=['GET', 'POST'])
def cit_form():
    if request.method == 'POST':
        data = {key: request.form.get(key, 0) for key in request.form}
        conn = mysql.connector.connect(**db_config)
        cursor = conn.cursor()

        query = """
            INSERT INTO cit (
                year, gross_total_income, deduction_80ia_business, deduction_sec37_disallowance,
                deduction_80g, net_taxable_income, tax_30_percent, tax_book_profit_18_5,
                tax_payable, surcharge_12, edu_cess_3, total_tax_payable, mat_credit,
                interest_234c, total_tax, advance_tax, tds, tcs, tax_on_assessment, refund
            ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
        """

        values = (
            data.get('year'),  # Include 'year' as the first value
            float(data.get('gross_total_income', 0)),
            float(data.get('deduction_80ia_business', 0)),
            float(data.get('deduction_sec37_disallowance', 0)),
            float(data.get('deduction_80g', 0)),
            float(data.get('net_taxable_income', 0)),
            float(data.get('tax_30_percent', 0)),
            float(data.get('tax_book_profit_18_5', 0)),
            float(data.get('tax_payable', 0)),
            float(data.get('surcharge_12', 0)),
            float(data.get('edu_cess_3', 0)),
            float(data.get('total_tax_payable', 0)),
            float(data.get('mat_credit', 0)),
            float(data.get('interest_234c', 0)),
            float(data.get('total_tax', 0)),
            float(data.get('advance_tax', 0)),
            float(data.get('tds', 0)),
            float(data.get('tcs', 0)),
            float(data.get('tax_on_assessment', 0)),
            float(data.get('refund', 0))
        )

        cursor.execute(query, values)
        conn.commit()
        cursor.close()
        conn.close()
        return redirect(url_for('index'))

    return render_template('cit_form.html')


@app.route('/itat', methods=['GET', 'POST'])
def itat_form():
    if request.method == 'POST':
        mat_tax_credit = request.form['mat_tax_credit']
        surcharge = request.form['surcharge']
        cess = request.form['cess']
        total_credit = request.form['total_credit']
        year=request.form['year']
        conn = mysql.connector.connect(**db_config)
        cursor = conn.cursor()
        cursor.execute("""
            INSERT INTO itat (year, mat_tax_credit, surcharge, cess, total_credit)
            VALUES (%s,%s, %s, %s, %s)
        """, (year,mat_tax_credit, surcharge, cess, total_credit))
        conn.commit()
        cursor.close()
        conn.close()
        return redirect(url_for('index'))
    return render_template('itat_form.html')


def get_db_connection():
    connection = mysql.connector.connect(**db_config)
    return connection

import pandas as pd
import pymysql
import io
@app.route('/reports')
def reports():
    return render_template("reports.html")



@app.route('/itr_report', methods=['GET'])
def itr_report():
    connection = pymysql.connect(**db_config)
    try:
        selected_year = request.args.get('year')

        if selected_year:
            # Fetch ITR data for the selected year
            query = "SELECT * FROM itr WHERE year = %s"
            df = pd.read_sql(query, connection, params=[selected_year])

            if df.empty:
                return "No records found for the selected year."

            # Transpose DataFrame: vertical fields, horizontal records
            df_transposed = df.transpose()
            df_transposed.insert(0, 'Field', df_transposed.index)

            # Rename columns as Record 1, Record 2, etc.
            record_cols = {i: f'Record {i}' for i in df_transposed.columns if isinstance(i, int)}
            df_transposed.rename(columns=record_cols, inplace=True)
            df_transposed.reset_index(drop=True, inplace=True)

            output = io.BytesIO()
            with pd.ExcelWriter(output, engine='xlsxwriter') as writer:
                df_transposed.to_excel(writer, index=False, sheet_name='ITR_Vertical')

                # Format for better readability (optional)
                workbook = writer.book
                worksheet = writer.sheets['ITR_Vertical']
                worksheet.set_column(0, 0, 30)  # Field column wider

            output.seek(0)
            return send_file(
                output,
                mimetype='application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
                as_attachment=True,
                download_name=f"ITR_Report_{selected_year}.xlsx"
            )

        else:
            # Render dropdown form with available years
            with connection.cursor() as cursor:
                cursor.execute("SELECT DISTINCT year FROM itr ORDER BY year DESC")
                years = [row[0] for row in cursor.fetchall()]
            return render_template("itr_reports.html", years=years)

    finally:
        connection.close()


@app.route('/ao_report', methods=['GET'])
def ao_report():
    selected_year = request.args.get('year')
    connection = pymysql.connect(**db_config)

    try:
        if selected_year:
            query = "SELECT * FROM ao WHERE year = %s"
            df = pd.read_sql(query, connection, params=[selected_year])

            if df.empty:
                return "No records found for the selected year."

            # Transpose the DataFrame: rows → fields, columns → records
            df_transposed = df.transpose()
            df_transposed.insert(0, 'Field', df_transposed.index)

            # Rename columns to "Record 1", "Record 2", ...
            for i in range(1, df_transposed.shape[1]):
                df_transposed.rename(columns={df_transposed.columns[i]: f"Record {i}"}, inplace=True)

            df_transposed.reset_index(drop=True, inplace=True)

            output = io.BytesIO()
            with pd.ExcelWriter(output, engine='xlsxwriter') as writer:
                df_transposed.to_excel(writer, index=False, sheet_name='AO_Vertical')

                # Optional: Adjust formatting
                workbook = writer.book
                worksheet = writer.sheets['AO_Vertical']
                worksheet.set_column(0, 0, 30)  # Widen 'Field' column

            output.seek(0)
            return send_file(
                output,
                mimetype='application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
                as_attachment=True,
                download_name=f"AO_Report_{selected_year}.xlsx"
            )
        else:
            with connection.cursor() as cursor:
                cursor.execute("SELECT DISTINCT year FROM ao ORDER BY year DESC")
                years = [row[0] for row in cursor.fetchall()]
            return render_template("ao_reports.html", years=years)
    finally:
        connection.close()


@app.route('/cit_report', methods=['GET'])
def cit_report():
    selected_year = request.args.get('year')
    connection = pymysql.connect(**db_config)

    try:
        if selected_year:
            # Fetch data from the `cit` table for the selected year
            query = "SELECT * FROM cit WHERE year = %s"
            df = pd.read_sql(query, connection, params=[selected_year])

            output = io.BytesIO()
            with pd.ExcelWriter(output, engine='xlsxwriter') as writer:
                workbook = writer.book

                # Write each row vertically on a separate sheet or below one another
                for i, (_, row) in enumerate(df.iterrows(), start=1):
                    # Convert the row to vertical format
                    vertical_df = pd.DataFrame(row).reset_index()
                    vertical_df.columns = ['Field', 'Value']

                    # Write each vertical entry below the previous (e.g., block by block)
                    start_row = (i - 1) * (len(vertical_df) + 3)  # 3-row gap between entries
                    vertical_df.to_excel(writer, sheet_name='CIT_Report', index=False, startrow=start_row)

            output.seek(0)

            return send_file(
                output,
                mimetype='application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
                as_attachment=True,
                download_name=f"CIT_Report_{selected_year}_Vertical.xlsx"
            )
        else:
            # Render dropdown for year selection
            with connection.cursor() as cursor:
                cursor.execute("SELECT DISTINCT year FROM cit ORDER BY year DESC")
                years = [row[0] for row in cursor.fetchall()]
            return render_template("cit_reports.html", years=years)
    finally:
        connection.close()

@app.route('/itat_report', methods=['GET'])
def itat_report():
    selected_year = request.args.get('year')
    connection = pymysql.connect(**db_config)

    try:
        if selected_year:
            query = "SELECT * FROM itat WHERE year = %s"
            df = pd.read_sql(query, connection, params=[selected_year])

            output = io.BytesIO()
            with pd.ExcelWriter(output, engine='xlsxwriter') as writer:
                df.T.to_excel(writer, header=False, sheet_name='ITAT_Report')
            output.seek(0)

            return send_file(
                output,
                mimetype='application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
                as_attachment=True,
                download_name=f"ITAT_Report_{selected_year}_Vertical.xlsx"
            )
        else:
            with connection.cursor() as cursor:
                cursor.execute("SELECT DISTINCT year FROM itat ORDER BY year DESC")
                years = [row[0] for row in cursor.fetchall()]
            return render_template("itat_reports.html", years=years)
    finally:
        connection.close()



@app.route('/itr_report_download', methods=['GET'])
def itr_report_download():
    connection = pymysql.connect(**db_config)
    try:
        selected_year = request.args.get('year')

        if selected_year:
            query = "SELECT * FROM itr WHERE year = %s"
            df = pd.read_sql(query, connection, params=[selected_year])

            output = io.BytesIO()
            with pd.ExcelWriter(output, engine='xlsxwriter') as writer:
                df.to_excel(writer, index=False, sheet_name=f"ITR {selected_year}")
            output.seek(0)

            return send_file(
                output,
                download_name=f"ITR_Report_{selected_year}.xlsx",
                as_attachment=True,
                mimetype='application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'
            )
        else:
            # If no year is selected, show dropdown
            with connection.cursor() as cursor:
                cursor.execute("SELECT DISTINCT year FROM itr ORDER BY year DESC")
                years = [row[0] for row in cursor.fetchall()]
            return render_template('itr_reports.html', years=years)
    finally:
        connection.close()


@app.route('/download/<int:doc_id>')
def download_report(doc_id):
    conn = get_db_connection()
    cursor = conn.cursor(dictionary=True)

    cursor.execute("SELECT * FROM documents WHERE id = %s", (doc_id,))
    document = cursor.fetchone()
    conn.close()

    if not document:
        return "Document not found", 404

    file_path = os.path.join('static', 'uploads', document['filename'])  # adjust as per your storage
    return send_from_directory(directory='static/uploads', path=document['filename'], as_attachment=True)

# main.py
from flask import Flask, send_file
import pandas as pd
import io
import pymysql  # or use mysql.connector if preferred
from config import db_config
@app.route('/summary_report', methods=['GET'])
def summary_report():
    year = request.args.get('year')

    if not year:
        connection = pymysql.connect(**db_config)
        try:
            years = set()
            for table in ['itr', 'ao', 'cit', 'itat']:
                df = pd.read_sql(f"SELECT DISTINCT year FROM {table}", connection)
                years.update(int(y) for y in df['year'].dropna().tolist())
            return render_template('summary_reports.html', years=sorted(years), message="Please select a year to download.")
        finally:
            connection.close()

    connection = pymysql.connect(**db_config)
    try:
        stages = ['itr', 'ao', 'cit', 'itat']
        stage_data = {}

        for stage in stages:
            query = f"SELECT * FROM {stage} WHERE year = %s"
            df = pd.read_sql(query, connection, params=[year])
            stage_data[stage.upper()] = df

        def safe_get(df, col):
            return df[col].values[0] if col in df.columns and not df.empty else '-'

        particulars = [
            "Gross Total Income", "Add: Disallowance u/s 14A", "Add: Disallowance u/s 37", "GTI as per",
            "Less: Deduction u/s 80IA", "Less: Deduction u/s 80G", "Net Taxable Income", "Tax @ 30%",
            "Tax @ 18.5% on Book Profit", "Surcharge @ 12%", "Education Cess @ 3%", "Total Tax Payable",
            "Less: MAT Credit", "Net Tax", "Add: Interest u/s 234C", "Total Tax",
            "Advance Tax", "TDS", "TCS", "SAT", "Tax on Regular Assessment", "Refund"
        ]

        columns = [
            'gross_total_income', 'disallowance_14a', 'disallowance_37', 'gti',
            'deduction_80ia', 'deduction_80g', 'net_taxable_income', 'tax_30',
            'book_profit_tax', 'surcharge_12', 'education_cess', 'total_tax',
            'mat_credit', 'net_tax', 'interest_234c', 'total_tax_payable',
            'advance_tax', 'tds', 'tcs', 'sat', 'tax_regular', 'refund'
        ]

        data = {
            "Particulars": particulars,
            "ITR": [safe_get(stage_data['ITR'], col) for col in columns],
            "AO": [safe_get(stage_data['AO'], col) for col in columns],
            "CIT(A)": [safe_get(stage_data['CIT'], col) for col in columns],
            "ITAT": [safe_get(stage_data['ITAT'], col) for col in columns],
        }

        df = pd.DataFrame(data)

        # Export to Excel with formatting
        output = io.BytesIO()
        with pd.ExcelWriter(output, engine='xlsxwriter') as writer:
            df.to_excel(writer, index=False, sheet_name=f'AY {year}')
            workbook = writer.book
            worksheet = writer.sheets[f'AY {year}']

            # Format definitions
            header_format = workbook.add_format({
                'bold': True,
                'text_wrap': True,
                'valign': 'middle',
                'align': 'center',
                'bg_color': '#007bff',
                'font_color': 'white',
                'border': 1
            })

            cell_format = workbook.add_format({
                'border': 1,
                'valign': 'top',
                'align': 'center',
            })

            # Apply formats
            for col_num, value in enumerate(df.columns):
                worksheet.write(0, col_num, value, header_format)
                # Auto column width
                max_len = max(df[value].astype(str).map(len).max(), len(str(value))) + 2
                worksheet.set_column(col_num, col_num, max_len)

            # Format data rows
            for row_num in range(1, len(df) + 1):
                for col_num in range(len(df.columns)):
                    worksheet.write(row_num, col_num, df.iloc[row_num - 1, col_num], cell_format)

        output.seek(0)

        return send_file(
            output,
            mimetype='application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
            as_attachment=True,
            download_name=f"Summary_Report_{year}.xlsx"
        )

    finally:
        connection.close()

if __name__ == '__main__':
      app.run(host='0.0.0.0', port=5003, debug=True)