import re

def fix_file(filepath):
    with open(filepath, 'r') as f:
        lines = f.readlines()

    new_lines = []
    i = 0
    fixed = 0
    total = len(lines)

    while i < total:
        line = lines[i]
        stripped = line.strip()

        if stripped == 'conn = None':
            new_lines.append(line)
            i += 1
            continue

        if not re.match(r'^\s*conn\s*=\s*get_db\s*\(\s*\)', line):
            new_lines.append(line)
            i += 1
            continue

        indent = len(line) - len(line.lstrip())
        indent_str = ' ' * indent

        # Check already wrapped (conn=None + try: before get_db)
        found_wrap = False
        for wi in range(max(0, i - 6), i):
            if lines[wi].strip() == 'conn = None':
                for tj in range(wi + 1, i):
                    if lines[tj].strip() == 'try:':
                        found_wrap = True
                        break
                if found_wrap:
                    break
        if found_wrap:
            new_lines.append(line)
            i += 1
            continue

        # ── Pattern B: try: immediately on next line ──
        if i + 1 < total and lines[i + 1].strip() == 'try:':
            try_indent = len(lines[i + 1]) - len(lines[i + 1].lstrip())
            finally_idx = None
            scan_end = min(i + 200, total)

            for sj in range(i + 2, scan_end):
                s = lines[sj].strip()
                if s == '':
                    continue
                sj_indent = len(lines[sj]) - len(lines[sj].lstrip())
                if sj_indent < try_indent:
                    break
                if sj_indent == try_indent:
                    if s == 'finally:':
                        finally_idx = sj
                        break
                    elif s.startswith('except'):
                        # After except, keep scanning for finally:
                        # (try/except/finally is valid)
                        continue
                    else:
                        # Other line at try_indent without finally — block ended
                        break

            if finally_idx is not None:
                # Find end of finally body
                finally_body_end = finally_idx + 1
                while (finally_body_end < total and
                       lines[finally_body_end].strip() != '' and
                       len(lines[finally_body_end]) - len(lines[finally_body_end].lstrip()) > try_indent):
                    finally_body_end += 1

                new_lines.append(f'{indent_str}conn = None\n')
                new_lines.append(f'{indent_str}try:\n')
                new_lines.append(f'{indent_str}    conn = get_db()\n')
                for bk in range(i + 2, finally_idx):
                    new_lines.append(lines[bk])
                new_lines.append(f'{indent_str}finally:\n')
                new_lines.append(f'{indent_str}    if conn:\n')
                new_lines.append(f'{indent_str}        conn.close()\n')

                i = finally_body_end
                fixed += 1
                continue
            else:
                # No finally — check for conn.close() after try/except
                close_idx = None
                for ci in range(i + 2, min(i + 200, total)):
                    s = lines[ci].strip()
                    if s and (s.startswith('def ') or s.startswith('class ')):
                        break
                    if s == 'conn.close()':
                        ci_indent = len(lines[ci]) - len(lines[ci].lstrip())
                        if ci_indent == indent:
                            close_idx = ci
                            break

                if close_idx is not None:
                    new_lines.append(f'{indent_str}conn = None\n')
                    new_lines.append(f'{indent_str}try:\n')
                    new_lines.append(f'{indent_str}    conn = get_db()\n')
                    for bk in range(i + 2, close_idx):
                        new_lines.append(lines[bk])
                    new_lines.append(f'{indent_str}finally:\n')
                    new_lines.append(f'{indent_str}    if conn:\n')
                    new_lines.append(f'{indent_str}        conn.close()\n')
                    i = close_idx + 1
                    fixed += 1
                    continue
                else:
                    new_lines.append(line)
                    i += 1
                    continue

        # ── Pattern A: find conn.close() at same indent ──
        close_idx = None
        for ci in range(i + 1, total):
            if lines[ci].strip() == 'conn.close()':
                ci_indent = len(lines[ci]) - len(lines[ci].lstrip())
                if ci_indent == indent:
                    close_idx = ci
                    break

        if close_idx is None:
            new_lines.append(line)
            i += 1
            continue

        # Check if there's an existing try/finally between get_db() and close()
        existing_try_idx = None
        existing_finally_idx = None
        for ti in range(i + 1, close_idx):
            if lines[ti].strip() == 'try:' and len(lines[ti]) - len(lines[ti].lstrip()) == indent:
                depth = 1
                fi = ti + 1
                while fi < close_idx and depth > 0:
                    s = lines[fi].strip()
                    if s == 'try:':
                        depth += 1
                    elif s.startswith('except') or s.startswith('finally'):
                        depth -= 1
                    fi += 1
                for fj in range(ti + 1, fi):
                    if lines[fj].strip() == 'finally:' and len(lines[fj]) - len(lines[fj].lstrip()) == indent:
                        for fc in range(fj + 1, close_idx + 1):
                            if lines[fc].strip() == 'conn.close()':
                                existing_try_idx = ti
                                existing_finally_idx = fj
                                break
                        if existing_try_idx:
                            break
                if existing_try_idx:
                    break

        if existing_try_idx is not None:
            # Existing try/finally — add conn=None, move get_db() inside try
            finally_body_end = existing_finally_idx + 1
            while (finally_body_end < total and
                   lines[finally_body_end].strip() != '' and
                   len(lines[finally_body_end]) - len(lines[finally_body_end].lstrip()) > indent):
                finally_body_end += 1

            new_lines.append(f'{indent_str}conn = None\n')
            new_lines.append(lines[existing_try_idx])
            for bk in range(i, existing_try_idx):
                if lines[bk].strip():
                    new_lines.append('    ' + lines[bk])
                else:
                    new_lines.append(lines[bk])
            for bk in range(existing_try_idx + 1, existing_finally_idx):
                new_lines.append(lines[bk])
            new_lines.append(lines[existing_finally_idx])
            new_lines.append(f'{indent_str}    if conn:\n')
            new_lines.append(f'{indent_str}        conn.close()\n')

            i = finally_body_end
            fixed += 1
        else:
            # No existing try/finally — simple wrap
            new_lines.append(f'{indent_str}conn = None\n')
            new_lines.append(f'{indent_str}try:\n')
            new_lines.append(f'{indent_str}    {stripped}\n')
            for bk in range(i + 1, close_idx):
                new_lines.append('    ' + lines[bk])
            new_lines.append(f'{indent_str}finally:\n')
            new_lines.append(f'{indent_str}    if conn:\n')
            new_lines.append(f'{indent_str}        conn.close()\n')

            i = close_idx + 1
            fixed += 1

    with open(filepath, 'w') as f:
        f.writelines(new_lines)
    return fixed

files = [
    'app/models.py',
    'app/routes/admin.py',
    'app/routes/agent_api.py',
    'app/routes/auth.py',
    'app/routes/email.py',
    'app/routes/invites.py',
    'app/routes/user_sites.py',
]

total_fixed = 0
for fp in files:
    count = fix_file(fp)
    if count:
        print(f'{fp}: {count} fixes')
    else:
        print(f'{fp}: 0 fixes')
    total_fixed += count
print(f'\nTotal: {total_fixed} fixes applied')