#!/usr/bin/env python3
"""
Applies "real progress" plumbing to your backend:

    python apply_progress_patch.py [directory containing the three files]

Edits:  ai_visibility.py, seo_analyzer.py, app.py
- Every anchor must match EXACTLY ONCE (whitespace-tolerant). If any anchor
  fails in a file, that file is left untouched and the failures are listed.
- A .bak copy is written before a file is modified.
- Safe to re-run: files that already contain the marker are skipped.
"""
import re
import sys
import shutil
import textwrap
from pathlib import Path

MARKER = "# --- live progress (real, event-driven) ---"


class Patcher:
    def __init__(self, path):
        self.path = Path(path)
        self.src = self.path.read_text(encoding="utf-8")
        self.errors = []

    def sub(self, name, old, new):
        """Replace a block that starts at the beginning of a line.
        `old` is matched token-by-token (whitespace-insensitive).
        `new` is written at indent 0 and re-indented to the anchor's indent."""
        toks = old.split()
        pat = re.compile(r"(?m)^([ \t]*)" + r"\s+".join(re.escape(t) for t in toks))
        ms = list(pat.finditer(self.src))
        if len(ms) != 1:
            self.errors.append(f"[{self.path.name}] '{name}': matched {len(ms)}x (expected 1)")
            return
        m = ms[0]
        indent = m.group(1)
        block = textwrap.dedent(new).strip("\n")
        block = "\n".join((indent + ln) if ln.strip() else ln for ln in block.split("\n"))
        self.src = self.src[: m.start()] + block + self.src[m.end():]

    def save(self):
        if self.errors:
            print(f"\n✗ {self.path.name}: NOT modified")
            for e in self.errors:
                print("   ", e)
            return False
        shutil.copy(self.path, str(self.path) + ".bak")
        self.path.write_text(self.src, encoding="utf-8")
        print(f"✓ {self.path.name}: patched (backup: {self.path.name}.bak)")
        return True


# ----------------------------------------------------------------------------
# ai_visibility.py
# ----------------------------------------------------------------------------
AI_HELPERS = f'''
{MARKER}
# Each stage owns a slice of the 0-100 bar. progress = slice_start + (done/total) * slice_size,
# so the bar only moves when real work completes.
def _stage_ranges(observed=True):
    if observed:
        return {{"discover": (0, 8), "crawl": (8, 40), "scan": (40, 60), "analyze": (60, 95), "report": (95, 100)}}
    return {{"discover": (0, 10), "crawl": (10, 50), "scan": (50, 90), "analyze": (90, 90), "report": (90, 100)}}


def _evt(stage, done=0, total=1, message="", observed=True, **extra):
    lo, hi = _stage_ranges(observed)[stage]
    frac = (done / total) if total else 1
    return {{"stage": stage, "done": done, "total": total,
            "progress": round(lo + (hi - lo) * frac, 1), "message": message, **extra}}


async def _emit(cb, evt):
    """Never let progress reporting break the analysis."""
    if cb is None:
        return
    try:
        r = cb(evt)
        if asyncio.iscoroutine(r):
            await r
    except Exception:
        pass


def _get_anthropic_client():'''


def patch_ai_visibility(path):
    p = Patcher(path)
    if MARKER in p.src:
        print(f"- {p.path.name}: already patched, skipping")
        return True

    p.sub("helpers", "def _get_anthropic_client():", AI_HELPERS)

    p.sub(
        "observed signature",
        "async def analyze_observed_ai_visibility(brand_name, target_prompts=None, competitor_names=None, use_live_ai=False, brand_facts=None, engines=None):",
        "async def analyze_observed_ai_visibility(brand_name, target_prompts=None, competitor_names=None, use_live_ai=False, brand_facts=None, engines=None, progress_cb=None):",
    )

    p.sub(
        "observed engine loop",
        """per_engine_raw = {}
    for engine in available:
        raw_results = []
        for prompt in target_prompts:
            answer = await loop.run_in_executor(None, _call_engine_sync, engine, prompt)
            raw_results.append({"prompt": prompt, "answer": answer})
        per_engine_raw[engine] = raw_results

    per_engine_metrics = {}
    for engine, raw_results in per_engine_raw.items():
        valid = [r for r in raw_results if r["answer"] is not None]
        representation = await _check_representation_accuracy(brand_facts, valid, loop)
        per_engine_metrics[engine] = _compute_engine_metrics(brand_name, target_prompts, competitor_names, raw_results, representation)""",
        """
        total_steps = len(available) * len(target_prompts) + len(available)  # prompts + one accuracy check per engine
        done = 0
        per_engine_raw = {}
        for engine in available:
            raw_results = []
            for prompt in target_prompts:
                answer = await loop.run_in_executor(None, _call_engine_sync, engine, prompt)
                raw_results.append({"prompt": prompt, "answer": answer})
                done += 1
                await _emit(progress_cb, _evt(
                    "analyze", done, total_steps,
                    f"Asking {ENGINE_DISPLAY_NAMES[engine]}",
                    engine=engine, answered=answer is not None,
                ))
            per_engine_raw[engine] = raw_results

        per_engine_metrics = {}
        for engine, raw_results in per_engine_raw.items():
            valid = [r for r in raw_results if r["answer"] is not None]
            representation = await _check_representation_accuracy(brand_facts, valid, loop)
            per_engine_metrics[engine] = _compute_engine_metrics(brand_name, target_prompts, competitor_names, raw_results, representation)
            done += 1
            await _emit(progress_cb, _evt("analyze", done, total_steps, f"Checking {ENGINE_DISPLAY_NAMES[engine]} answers"))
        """,
    )

    p.sub("run signature", "include_raw_metrics=False,", "include_raw_metrics=False,\nprogress_cb=None,")

    p.sub(
        "discover + crawl",
        """urls = await discover_urls_parallel(base_url, max_pages)
    if not urls:
        urls = [base_url]

    playwright_data, failed_fetches = await fetch_all_pages_parallel(urls, max_concurrent)""",
        """
        obs = bool(auto_measure_observed)

        await _emit(progress_cb, _evt("discover", 0, 1, "Reading sitemap & internal links", obs))
        urls = await discover_urls_parallel(base_url, max_pages)
        if not urls:
            urls = [base_url]
        await _emit(progress_cb, _evt("discover", 1, 1, f"Found {len(urls)} pages", obs, urls_found=len(urls)))

        async def _on_page(done, total, page):
            await _emit(progress_cb, _evt(
                "crawl", done, total, "Crawling pages", obs,
                url=page.get("requested_url") or page.get("url"),
                ok=not page.get("fetch_failed"),
            ))

        playwright_data, failed_fetches = await fetch_all_pages_parallel(urls, max_concurrent, on_page=_on_page)
        """,
    )

    p.sub(
        "scan",
        """robots_txt = await _fetch_robots_txt(base_url)
    robots_bot_permissions = _parse_robots_txt(robots_txt, path="/") if robots_txt else None
    sitemap_present, llms_txt_present = await _check_sitemap_and_llms_txt(base_url)

    page_results = await asyncio.gather(*[
        analyze_page_ai_visibility(
            seo_data, domain, use_ai=use_ai,
            robots_bot_permissions=robots_bot_permissions,
            sitemap_present=sitemap_present, llms_txt_present=llms_txt_present,
        )
        for seo_data in playwright_data
    ])""",
        """
        scan_total = len(playwright_data) + 2  # pages + robots.txt + sitemap/llms.txt
        await _emit(progress_cb, _evt("scan", 0, scan_total, "Checking robots.txt", obs))
        robots_txt = await _fetch_robots_txt(base_url)
        robots_bot_permissions = _parse_robots_txt(robots_txt, path="/") if robots_txt else None
        await _emit(progress_cb, _evt("scan", 1, scan_total, "Checking sitemap & llms.txt", obs))
        sitemap_present, llms_txt_present = await _check_sitemap_and_llms_txt(base_url)
        await _emit(progress_cb, _evt("scan", 2, scan_total, "Scanning page content", obs))

        scanned = 2

        async def _analyze_one(seo_data):
            nonlocal scanned
            res = await analyze_page_ai_visibility(
                seo_data, domain, use_ai=use_ai,
                robots_bot_permissions=robots_bot_permissions,
                sitemap_present=sitemap_present, llms_txt_present=llms_txt_present,
            )
            scanned += 1
            await _emit(progress_cb, _evt("scan", scanned, scan_total, "Scanning page content", obs, url=seo_data.get("url")))
            return res

        page_results = await asyncio.gather(*[_analyze_one(d) for d in playwright_data])
        """,
    )

    p.sub(
        "auto-generate emit",
        "if auto_measure_observed:",
        """
        if auto_measure_observed:
            await _emit(progress_cb, _evt("analyze", 0, 1, "Inferring customer prompts from your homepage", True))
        """,
    )
    # NOTE: the patch above replaces only the `if` line; the original body that
    # follows stays attached to it (the emit is the first statement of the block).

    p.sub(
        "observed call",
        """observed = await analyze_observed_ai_visibility(
        brand_name=inferred_brand_name or domain, target_prompts=target_prompts,
        competitor_names=competitor_names, use_live_ai=use_live_ai_for_observed,
        brand_facts=brand_facts,
    )""",
        """
        observed = await analyze_observed_ai_visibility(
            brand_name=inferred_brand_name or domain, target_prompts=target_prompts,
            competitor_names=competitor_names, use_live_ai=use_live_ai_for_observed,
            brand_facts=brand_facts, progress_cb=progress_cb,
        )
        await _emit(progress_cb, _evt("report", 0, 1, "Building report", obs))
        """,
    )
    return p.save()


# ----------------------------------------------------------------------------
# seo_analyzer.py
# ----------------------------------------------------------------------------
SEO_HELPERS = f'''
{MARKER}
_SEO_STAGE_RANGES = {{"discover": (0, 8), "crawl": (8, 45), "scan": (45, 70), "analyze": (70, 95), "report": (95, 100)}}


def _evt(stage, done=0, total=1, message="", **extra):
    lo, hi = _SEO_STAGE_RANGES[stage]
    frac = (done / total) if total else 1
    return {{"stage": stage, "done": done, "total": total,
            "progress": round(lo + (hi - lo) * frac, 1), "message": message, **extra}}


async def _emit(cb, evt):
    """Never let progress reporting break the analysis."""
    if cb is None:
        return
    try:
        r = cb(evt)
        if asyncio.iscoroutine(r):
            await r
    except Exception:
        pass


async def fetch_robots_txt(base_url):'''


def patch_seo_analyzer(path):
    p = Patcher(path)
    if MARKER in p.src:
        print(f"- {p.path.name}: already patched, skipping")
        return True

    p.sub("helpers", "async def fetch_robots_txt(base_url):", SEO_HELPERS)

    # latent bug: callers unpack two values, this returned a bare list
    p.sub("empty-return fix", 'print("[fail] No URLs to fetch!")\n        return []',
          'print("[fail] No URLs to fetch!")\nreturn [], []')

    p.sub("fetch signature", "async def fetch_all_pages_parallel(urls, max_concurrent=1):",
          "async def fetch_all_pages_parallel(urls, max_concurrent=1, on_page=None):")

    p.sub(
        "fetch loop hook",
        "result = await task\n            results.append(result)",
        """
        result = await task
        results.append(result)
        if on_page and result is not None:
            try:
                r = on_page(i + 1, len(urls), result)
                if asyncio.iscoroutine(r):
                    await r
            except Exception:
                pass
        """,
    )

    p.sub("suggestions signature", "async def generate_all_page_suggestions_parallel(pages):",
          "async def generate_all_page_suggestions_parallel(pages, progress_cb=None):")

    p.sub(
        "suggestions tracker",
        'suggestions = {}\n\n    print(f"[ai] Generating AI suggestions for {len(pages)} pages in parallel...")',
        """
        suggestions = {}

        print(f"[ai] Generating AI suggestions for {len(pages)} pages in parallel...")

        _ai_done = {"n": 0}

        async def _tracked(coro):
            try:
                return await coro
            finally:
                _ai_done["n"] += 1
                await _emit(progress_cb, _evt("analyze", _ai_done["n"], len(pages), "Generating AI suggestions"))
        """,
    )
    p.sub("suggestions wrap", "task = generate_page_suggestions_async(page)",
          "task = _tracked(generate_page_suggestions_async(page))")
    p.sub("add_comprehensive signature", "async def add_comprehensive_suggestions_async(results):",
          "async def add_comprehensive_suggestions_async(results, progress_cb=None):")
    p.sub("add_comprehensive call", "suggestions_dict = await generate_all_page_suggestions_parallel(results)",
          "suggestions_dict = await generate_all_page_suggestions_parallel(results, progress_cb=progress_cb)")

    p.sub(
        "orchestrator signature",
        'async def run_seo_and_suggestions_async(base_url, max_pages=20, tmp_dir="/tmp", use_ai=True, max_concurrent=3):',
        'async def run_seo_and_suggestions_async(base_url, max_pages=20, tmp_dir="/tmp", use_ai=True, max_concurrent=3, progress_cb=None):',
    )
    p.sub("emit discover start", "# STEP 1: Discover URLs in parallel",
          '# STEP 1: Discover URLs in parallel\nawait _emit(progress_cb, _evt("discover", 0, 1, "Reading sitemap & internal links"))')
    p.sub("emit discover done", 'print(f"[info] Found {len(urls)} URLs to analyze")',
          'print(f"[info] Found {len(urls)} URLs to analyze")\n'
          'await _emit(progress_cb, _evt("discover", 1, 1, f"Found {len(urls)} pages", urls_found=len(urls)))')
    p.sub(
        "crawl hook",
        "playwright_data, failed_fetches = await fetch_all_pages_parallel(urls, max_concurrent)",
        """
        async def _on_page(done, total, page):
            await _emit(progress_cb, _evt(
                "crawl", done, total, "Crawling pages",
                url=page.get("requested_url") or page.get("url"),
                ok=not page.get("fetch_failed"),
            ))

        playwright_data, failed_fetches = await fetch_all_pages_parallel(urls, max_concurrent, on_page=_on_page)
        """,
    )
    p.sub(
        "scan emit",
        "results = await analyze_pages_parallel(playwright_data, domain)",
        """
        await _emit(progress_cb, _evt("scan", 0, 1, "Analyzing page content"))
        results = await analyze_pages_parallel(playwright_data, domain)
        await _emit(progress_cb, _evt("scan", 1, 1, "Page content analyzed"))
        """,
    )
    p.sub("ai suggestions call", "await add_comprehensive_suggestions_async(results)",
          "await add_comprehensive_suggestions_async(results, progress_cb=progress_cb)")
    p.sub("report emit", "# STEP 6: Save to CSV",
          '# STEP 6: Save to CSV\nawait _emit(progress_cb, _evt("report", 0, 1, "Building report"))')

    p.sub("full audit signature",
          "async def run_full_seo_audit(base_url, max_pages=20, use_ai=True, max_concurrent=3):",
          "async def run_full_seo_audit(base_url, max_pages=20, use_ai=True, max_concurrent=3, progress_cb=None):")
    p.sub("full audit call",
          'tmp_dir="/tmp",\n        use_ai=use_ai,\n        max_concurrent=max_concurrent,',
          'tmp_dir="/tmp",\nuse_ai=use_ai,\nmax_concurrent=max_concurrent,\nprogress_cb=progress_cb,')
    return p.save()


# ----------------------------------------------------------------------------
# app.py
# ----------------------------------------------------------------------------
APP_HELPERS = f'''
{MARKER}
async def _stream_job(make_coro):
    """Run an analysis coroutine and yield its progress events as they happen,
    then the final result. `make_coro(cb)` must return the awaitable."""
    q: asyncio.Queue = asyncio.Queue()

    async def cb(evt):
        await q.put(evt)

    task = asyncio.create_task(make_coro(cb))
    try:
        while not task.done() or not q.empty():
            try:
                evt = await asyncio.wait_for(q.get(), timeout=0.5)
            except asyncio.TimeoutError:
                continue
            yield {{"type": "progress", **evt}}
        result = await task
        if isinstance(result, dict) and result.get("status") == "error":
            yield {{"type": "error", "message": result.get("message", "Analysis failed")}}
        else:
            yield {{"type": "result", "data": json.loads(json.dumps(result, default=str))}}
    except Exception as e:
        yield {{"type": "error", "message": str(e)}}
    finally:
        if not task.done():
            task.cancel()


def _normalize_url(url):
    return url if str(url).startswith(("http://", "https://")) else f"https://{{url}}"


async def analyze_ai_stream(url, max_pages, max_concurrent, use_ai, measure_observed):
    async for msg in _stream_job(lambda cb: run_ai_visibility_analysis(
        base_url=_normalize_url(url), max_pages=int(max_pages), max_concurrent=int(max_concurrent),
        use_ai=use_ai, auto_measure_observed=bool(measure_observed), progress_cb=cb,
    )):
        yield msg


async def analyze_seo_stream(url, max_pages, max_concurrent, use_ai):
    async for msg in _stream_job(lambda cb: run_full_seo_audit(
        base_url=_normalize_url(url), max_pages=int(max_pages),
        use_ai=use_ai, max_concurrent=int(max_concurrent), progress_cb=cb,
    )):
        yield msg


# ---- Gradio Interface ----'''

APP_WIRING = '''
# Streaming endpoints (hidden; consumed by the frontend via @gradio/client)
with demo:
    ai_stream_out = gr.JSON(visible=False)
    ai_stream_btn = gr.Button("stream-ai", visible=False)
    ai_stream_btn.click(
        fn=analyze_ai_stream,
        inputs=[url_input, max_pages_input, max_concurrent_input, use_ai_check, measure_observed_check],
        outputs=ai_stream_out,
        api_name="analyze_ai_stream",
    )
    seo_stream_out = gr.JSON(visible=False)
    seo_stream_btn = gr.Button("stream-seo", visible=False)
    seo_stream_btn.click(
        fn=analyze_seo_stream,
        inputs=[url_input, max_pages_input, max_concurrent_input, use_ai_check],
        outputs=seo_stream_out,
        api_name="analyze_seo_stream",
    )

if __name__ == "__main__":'''


def patch_app(path):
    p = Patcher(path)
    if MARKER in p.src:
        print(f"- {p.path.name}: already patched, skipping")
        return True
    p.sub("json import", "import gradio as gr", "import gradio as gr\nimport json")
    p.sub("stream helpers", "# ---- Gradio Interface ----", APP_HELPERS)
    p.sub("stream wiring", 'if __name__ == "__main__":', APP_WIRING)
    return p.save()


if __name__ == "__main__":
    root = Path(sys.argv[1] if len(sys.argv) > 1 else ".")
    ok = True
    for fn, name in ((patch_ai_visibility, "ai_visibility.py"),
                     (patch_seo_analyzer, "seo_analyzer.py"),
                     (patch_app, "app.py")):
        f = root / name
        if not f.exists():
            print(f"✗ {name}: not found in {root.resolve()}")
            ok = False
            continue
        ok = fn(f) and ok
    sys.exit(0 if ok else 1)