#!/usr/bin/env python3
"""
Fix relative image URLs in gov_raw content.
For each record, parse all <img src="..."> relative URLs
and resolve them against the record's page_url (or source_url).
"""
import re
import sqlite3
import os
from urllib.parse import urljoin, urlparse

DB_PATH = os.getenv("SEARCH_DB", "/root/search.db")

def fix_content_images(content, base_url):
    """Fix relative image src URLs in content using base_url."""
    if not content or not base_url:
        return content

    def fix_src(match):
        tag = match.group(0)
        src_match = re.search(r'src="([^"]*)"', tag)
        if not src_match:
            return tag
        src = src_match.group(1)
        # If already absolute, skip
        if src.startswith('http://') or src.startswith('https://') or src.startswith('data:'):
            return tag
        # Only fix meaningful images (skip tiny file-type icons)
        if '/filetype/' in src:
            return tag
        if '/icon16/' in src:
            return tag
        # Resolve relative URL against base_url
        absolute = urljoin(base_url, src)
        new_tag = tag.replace('src="' + src + '"', 'src="' + absolute + '"')
        return new_tag

    return re.sub(r'<img[^>]*>', fix_src, content, flags=re.IGNORECASE)

def main():
    db = sqlite3.connect(DB_PATH)
    db.row_factory = sqlite3.Row
    c = db.cursor()

    # Count total
    total = c.execute('SELECT COUNT(*) FROM gov_raw WHERE content LIKE "%<img%"').fetchone()[0]
    print(f'Total records with <img> tags: {total}')

    # Process in batches
    batch_size = 500
    offset = 0
    fixed = 0
    skipped = 0
    errors = 0

    while offset < total:
        rows = c.execute(
            'SELECT id, content, page_url, source_url, site_name FROM gov_raw WHERE content LIKE "%<img%" ORDER BY id LIMIT ? OFFSET ?',
            (batch_size, offset)
        ).fetchall()

        for row in rows:
            content = row['content']
            base_url = row['page_url'] or row['source_url'] or ''
            if not base_url:
                skipped += 1
                continue

            new_content = fix_content_images(content, base_url)
            if new_content != content:
                try:
                    c.execute('UPDATE gov_raw SET content=? WHERE id=?', (new_content, row['id']))
                    fixed += 1
                except Exception as e:
                    print(f'  ERROR id={row["id"]}: {e}')
                    errors += 1

        offset += batch_size
        db.commit()
        print(f'  Progress: {offset}/{total}, fixed={fixed}, skipped={skipped}, errors={errors}')

    db.close()
    print(f'\nDone! Fixed: {fixed}, Skipped: {skipped}, Errors: {errors}')

if __name__ == '__main__':
    main()