from pathlib import Path
p = Path('/root/homelab-brain/homelab-ai-bot/tools/deep_research.py')
text = p.read_text()

# 1. Add datetime import
text = text.replace(
    'import logging\nimport re\nimport time',
    'import logging\nimport re\nimport time\nfrom datetime import date'
)

# 2. Add _inject_date_context function before _price_report_quality
text = text.replace(
    'def _price_report_quality(report: str):',
    '''def _inject_date_context(query):
    today = date.today()
    return (
        f"Heutiges Datum: {today.strftime('%d. %B %Y')}. "
        f"'Letzte 3 Monate' = {(today.replace(day=1) - __import__('datetime').timedelta(days=90)).strftime('%B %Y')} bis {today.strftime('%B %Y')}.\n\n"
        + query
    )


def _price_report_quality(report: str):'''
)

# 3. Add year check in _price_report_quality before return
text = text.replace(
    '    return len(missing) == 0, missing',
    '''    today = date.today()
    ok_years = {str(today.year), str(today.year - 1)}
    found = set(re.findall(r"20\d{2}", text))
    if found and not (found & ok_years):
        missing.append(f"Jahre falsch: {found} statt {ok_years}")

    return len(missing) == 0, missing'''
)

# 4. Inject date context in handle_deep_research
text = text.replace(
    '    log.info("deep_research gestartet: %s", query[:120])',
    '    query = _inject_date_context(query)\n    log.info("deep_research gestartet: %s", query[:200])'
)

p.write_text(text)
print('OK')