app.py

"""
app.py — CodeModernizer Web UI Backend (Fixed for Kaggle & Context)
"""

import asyncio
import json
import os
import shutil
import subprocess
import sys
import threading
import time
import urllib.request
import zipfile
from collections import deque
from pathlib import Path
from typing import Optional

from dotenv import load_dotenv
load_dotenv()

from google import genai

import uvicorn
from fastapi import FastAPI, File, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse, StreamingResponse
from pydantic import BaseModel

app = FastAPI(title="CodeModernizer UI")
app.add_middleware(
    CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"]
)

# Works in both .py scripts and Kaggle/Jupyter notebooks
try:
    PROJECT_ROOT = Path(__file__).parent.resolve()
except NameError:
    # Notebook environment — app.py lives in /kaggle/working/project/
    PROJECT_ROOT = Path("/kaggle/working/project")

# ─── Fixed directory paths ────────────────────────────────────────────────────
FIXED_DIRS = {
    "source": PROJECT_ROOT / "uploaded_repo",
    "phase1": PROJECT_ROOT / "phase1_output",
    "phase2": PROJECT_ROOT / "phase2_storage",
    "phase3": PROJECT_ROOT / "phase3_output",
}

GEMMA_API_KEY = os.environ.get("GEMMA_API_KEY", "")
_genai_client = genai.Client(api_key=GEMMA_API_KEY)

state: dict = {
    "repo_path":        None,
    "repo_is_managed":  False,
    "phase1_output":    str(FIXED_DIRS["phase1"]),
    "transform_output": str(FIXED_DIRS["phase3"]),
    "phase2_storage":   str(FIXED_DIRS["phase2"]),
    "transform_mode":   None,
    "running":          False,
    "current_phase":    None,
    "log_buffer":       deque(maxlen=1000),
    "phase1_done":      False,
    "transform_done":   False,
    "gemma_model":      "gemma-4-31b-it",
}

class GithubRequest(BaseModel):
    url: str

class ConfigRequest(BaseModel):
    repo_path:        Optional[str] = None
    phase1_output:    Optional[str] = None
    transform_output: Optional[str] = None

class TransformRequest(BaseModel):
    mode:         str
    model:        str  = "gemma-4-31b-it"
    retry_failed: bool = False

class QueryRequest(BaseModel):
    question:   str
    orig_file:  Optional[str] = None
    trans_file: Optional[str] = None
    mode:       str            = "auto"

class ResetRequest(BaseModel):
    targets: list[str]   # ["source", "phase1", "phase3"]


def _log(msg: str):
    ts   = time.strftime("%H:%M:%S")
    line = f"[{ts}] {msg}"
    state["log_buffer"].append(line)
    print(line, flush=True)


def _run_subprocess(cmd: list[str], cwd: str | None = None) -> int:
    env = os.environ.copy()
    env["PYTHONUNBUFFERED"] = "1"
    proc = subprocess.Popen(
        cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
        text=True, encoding="utf-8", errors="replace",
        cwd=cwd or str(PROJECT_ROOT), env=env,
    )
    for raw in iter(proc.stdout.readline, ""):
        if (line := raw.rstrip()):
            state["log_buffer"].append(line)
            print(line, flush=True)
    proc.wait()
    return proc.returncode


def _gemma_is_running() -> bool:
    """Gemma API is a cloud service — always available if API key is valid."""
    try:
        _genai_client.models.generate_content(
            model=state.get("gemma_model", "gemma-4-31b-it"),
            contents="ping",
        )
        return True
    except Exception:
        return False


def _gemma_call_sync(prompt: str, model: str, system: str = "") -> str:
    full_prompt = f"{system}\n\n{prompt}" if system else prompt
    try:
        response = _genai_client.models.generate_content(
            model=model,
            contents=full_prompt,
        )
        return response.text.strip()
    except Exception as e:
        raise RuntimeError(f"Gemma API call failed: {e}")


async def _gemma_call(prompt: str, model: str, system: str = "") -> str:
    loop = asyncio.get_event_loop()
    return await loop.run_in_executor(None, lambda: _gemma_call_sync(prompt, model, system))


def _refresh_done_flags():
    state["phase1_done"]    = Path(state["phase1_output"], "ir_registry.json").exists()
    state["transform_done"] = (
        Path(state["transform_output"]).exists() and
        any(Path(state["transform_output"]).iterdir())
        if Path(state["transform_output"]).exists() else False
    )


def _restore_state_from_disk():
    # Always use the fixed uploaded_repo path
    if state["repo_path"] is None:
        uploaded = FIXED_DIRS["source"]
        if uploaded.exists() and list(uploaded.rglob("*.java")):
            children = [c for c in uploaded.iterdir()
                        if c.is_dir() and not c.name.startswith("__")]
            if len(children) == 1 and any(children[0].rglob("*.java")):
                state["repo_path"] = str(children[0])
            else:
                state["repo_path"] = str(uploaded)
            state["repo_is_managed"] = True
            print(f"[RESTORE] repo_path ← {state['repo_path']}")

    _refresh_done_flags()
    if state["phase1_done"]:
        print(f"[RESTORE] phase1_done ← {state['phase1_output']}")
    if state["transform_done"]:
        print(f"[RESTORE] transform_done ← {state['transform_output']}")


@app.on_event("startup")
async def on_startup():
    _restore_state_from_disk()


@app.get("/api/config")
async def get_config():
    return {k: state[k] for k in
            ("repo_path","phase1_output","transform_output","phase2_storage","gemma_model")}

@app.post("/api/config")
async def set_config(req: ConfigRequest):
    if req.repo_path:
        state["repo_path"]       = req.repo_path
        state["repo_is_managed"] = False
    if req.phase1_output:    state["phase1_output"]    = req.phase1_output
    if req.transform_output: state["transform_output"] = req.transform_output
    _refresh_done_flags()
    return {"status": "ok"}


def _dir_info(path: Path) -> dict:
    if not path.exists():
        return {"exists": False, "size_mb": 0.0, "files": 0, "path": str(path)}
    all_files = [f for f in path.rglob("*") if f.is_file()]
    size = sum(f.stat().st_size for f in all_files)
    return {"exists": True, "size_mb": round(size / 1_048_576, 2),
            "files": len(all_files), "path": str(path)}


# ─── Improved Reset Logic ─────────────────────────────────────────────────────
@app.post("/api/reset")
async def reset_outputs(req: ResetRequest):
    if state["running"]:
        raise HTTPException(409, "Cannot reset while a job is running.")

    removed: list[str] = []
    errors:  list[str] = []

    # Cascade clear Phase 2 if Phase 1 is cleared
    targets = set(req.targets)
    if "phase1" in targets:
        targets.add("phase2")

    for target in targets:
        # Resolve target path (handle manual source overrides)
        if target == "source" and state["repo_path"] and not state["repo_is_managed"]:
            path = Path(state["repo_path"])
        else:
            path = FIXED_DIRS.get(target)
            
        if not path:
            continue

        try:
            if path.exists():
                shutil.rmtree(path)
                # Ensure the folder is recreated so future ops don't hit "Not Found"
                path.mkdir(parents=True, exist_ok=True)
                _log(f"[RESET] Cleared and recreated: {path}")
                removed.append(str(path))
            else:
                removed.append(f"{path} (already empty)")

            # Clear state dependencies
            if target == "source":
                state["repo_path"]       = None
                state["repo_is_managed"] = False
            elif target == "phase1":
                state["phase1_done"] = False
            elif target == "phase3":
                state["transform_done"] = False

        except Exception as e:
            errors.append(f"{target}: {e}")

    _refresh_done_flags()
    if errors:
        raise HTTPException(500, "; ".join(errors))
    return {"removed": removed, "status": "ok"}


@app.get("/api/reset/info")
async def reset_info():
    return {k: _dir_info(v) for k, v in FIXED_DIRS.items()}


@app.get("/api/debug/paths")
async def debug_paths():
    """Call this to verify PROJECT_ROOT resolved correctly on your platform."""
    return {
        "PROJECT_ROOT": str(PROJECT_ROOT),
        "dirs": {k: {"path": str(v), "exists": v.exists()} for k, v in FIXED_DIRS.items()},
    }


@app.get("/api/gemma/status")
async def gemma_status():
    return {"running": True, "model": state["gemma_model"]}


# ─── Source setup ─────────────────────────────────────────────────────────────
@app.post("/api/setup/upload")
async def upload_zip(file: UploadFile = File(...)):
    if not file.filename.lower().endswith(".zip"):
        raise HTTPException(400, "Only .zip files are accepted")

    dest_dir = PROJECT_ROOT / "uploaded_repo"
    if dest_dir.exists():
        shutil.rmtree(dest_dir)
    dest_dir.mkdir(parents=True)

    zip_path = PROJECT_ROOT / "_upload.zip"
    zip_path.write_bytes(await file.read())
    with zipfile.ZipFile(zip_path) as zf:
        zf.extractall(dest_dir)
    zip_path.unlink(missing_ok=True)

    children = [c for c in dest_dir.iterdir() if not c.name.startswith("__")]
    actual_root = dest_dir
    if len(children) == 1 and children[0].is_dir():
        actual_root = children[0]

    state["repo_path"]       = str(actual_root)
    state["repo_is_managed"] = True
    java_count = len(list(actual_root.rglob("*.java")))
    return {"path": str(actual_root), "java_files": java_count, "status": "ok"}


@app.post("/api/setup/github")
async def clone_github(req: GithubRequest):
    if state["running"]:
        raise HTTPException(409, "A job is already running")

    clone_dir = PROJECT_ROOT / "github_repo"
    state.update(running=True, current_phase="clone")
    state["log_buffer"].clear()

    def _clone():
        try:
            if clone_dir.exists(): shutil.rmtree(clone_dir)
            _log(f"Cloning {req.url} …")
            rc = _run_subprocess(["git", "clone", "--depth=1", req.url, str(clone_dir)])
            if rc == 0:
                state["repo_path"]       = str(clone_dir)
                state["repo_is_managed"] = True
                _log(f"Clone complete ✓  ({len(list(clone_dir.rglob('*.java')))} Java files)")
            else:
                _log("ERROR: git clone failed")
        finally:
            state["running"]       = False
            state["current_phase"] = None

    threading.Thread(target=_clone, daemon=True).start()
    return {"status": "cloning"}


# ─── Phase 1 ──────────────────────────────────────────────────────────────────
@app.post("/api/phase1/run")
async def run_phase1():
    if state["running"]:       raise HTTPException(409, "A job is already running")
    if not state["repo_path"]: raise HTTPException(400, "No repository source configured")

    state.update(running=True, current_phase="phase1", phase1_done=False)
    state["log_buffer"].clear()

    def _work():
        try:
            _log(f"Phase 1 — IR Extraction")
            _log(f"Repo   : {state['repo_path']}")
            _log(f"Output : {state['phase1_output']}")
            _log("─" * 50)
            sys.path.insert(0, str(PROJECT_ROOT))
            from engine.phase1_runner import run_phase1 as _p1
            _p1(repo_path=state["repo_path"], output_dir=state["phase1_output"])
            state["phase1_done"] = True
            _log("─" * 50)
            _log("Phase 1 complete ✓")
        except Exception as exc:
            _log(f"ERROR: {exc}")
        finally:
            state["running"]       = False
            state["current_phase"] = None

    threading.Thread(target=_work, daemon=True).start()
    return {"status": "started"}


# ─── Phase 3 Transform ────────────────────────────────────────────────────────
@app.post("/api/transform/run")
async def run_transform(req: TransformRequest):
    if state["running"]:       raise HTTPException(409, "A job is already running")
    if not state["repo_path"]: raise HTTPException(400, "No repository source configured")

    state.update(running=True, current_phase="transform",
                 transform_mode=req.mode, transform_done=False)
    state["log_buffer"].clear()
    state["gemma_model"] = req.model

    def _work():
        try:
            _log(f"Phase 3 -- {'Java -> Java 21' if req.mode == 'java21' else 'Java -> Python'}")
            _log(f"Model  : {req.model}")
            _log("─" * 50)
            cmd = [
                sys.executable, str(PROJECT_ROOT / "run_phase3.py"),
                "--repo",   state["repo_path"],
                "--output", state["transform_output"],
                "--phase1", state["phase1_output"],
                "--model",  req.model,
            ]
            if req.mode == "python": cmd += ["--source", "java", "--target", "python"]
            else:                    cmd += ["--target-java", "21"]
            if req.retry_failed:     cmd += ["--retry-failed"]

            rc = _run_subprocess(cmd, cwd=str(PROJECT_ROOT))
            state["transform_done"] = True
            _log("─" * 50)
            _log(f"Transform complete (exit {rc}) ✓" if rc == 0 else f"Transform finished (exit {rc})")
        except Exception as exc:
            _log(f"ERROR: {exc}")
        finally:
            state["running"]       = False
            state["current_phase"] = None

    threading.Thread(target=_work, daemon=True).start()
    return {"status": "started"}


@app.post("/api/transform/stop")
async def stop_transform():
    state["running"]       = False
    state["current_phase"] = None
    return {"status": "stop_requested"}


@app.get("/api/transform/summary")
async def transform_summary():
    out = Path(state["transform_output"])
    summary: dict = {}
    for sf in (out.glob("_status*.json") if out.exists() else []):
        try:
            for entry in json.loads(sf.read_text()).values():
                s = entry.get("status", "unknown")
                summary[s] = summary.get(s, 0) + 1
        except Exception:
            pass
    return {"summary": summary, "output_dir": str(out)}


# ─── Logs (SSE) ───────────────────────────────────────────────────────────────
@app.get("/api/logs/stream")
async def stream_logs():
    async def _gen():
        sent = 0
        while True:
            buf = list(state["log_buffer"])
            if len(buf) > sent:
                for line in buf[sent:]:
                    yield f"data: {json.dumps({'line': line})}\n\n"
                sent = len(buf)
            if not state["running"]:
                yield f"data: {json.dumps({'done': True, 'sent': sent})}\n\n"
                return
            await asyncio.sleep(0.25)

    return StreamingResponse(_gen(), media_type="text/event-stream",
                             headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"})


@app.get("/api/status")
async def get_status():
    _refresh_done_flags()
    return {
        "running":          state["running"],
        "current_phase":    state["current_phase"],
        "phase1_done":      state["phase1_done"],
        "transform_done":   state["transform_done"],
        "repo_path":        state["repo_path"],
        "repo_is_managed":  state["repo_is_managed"],
        "transform_mode":   state["transform_mode"],
        "gemma_running":    True,
    }


# ─── File browser ─────────────────────────────────────────────────────────────
IGNORED = {"target","build",".git","__pycache__","node_modules","bin","out",".idea"}

@app.get("/api/files")
async def list_files(kind: str = "original"):
    if kind == "original":
        root = Path(state["repo_path"]) if state["repo_path"] else None
        exts = {".java"}
    else:
        root = Path(state["transform_output"])
        exts = {".py"} if root and root.exists() and any(root.rglob("*.py")) else {".java"}

    if not root or not root.exists():
        return {"files": [], "root": None}

    files = []
    for ext in exts:
        for p in root.rglob(f"*{ext}"):
            if any(part in IGNORED for part in p.parts):
                continue
            files.append(str(p.relative_to(root)))

    return {"files": sorted(files), "root": str(root)}


@app.get("/api/file")
async def get_file(path: str, kind: str = "original"):
    root = (Path(state["repo_path"]) if kind == "original" and state["repo_path"]
            else Path(state["transform_output"]))

    if not root or not root.exists():
        raise HTTPException(404, "Root directory not found")

    candidate = root / path
    if not candidate.exists():
        swapped = candidate.with_suffix(".py" if candidate.suffix == ".java" else ".java")
        candidate = swapped if swapped.exists() else None

    if not candidate or not candidate.exists():
        raise HTTPException(404, f"File not found: {path}")

    try:
        return {"content": candidate.read_text(encoding="utf-8", errors="replace"),
                "path":    str(candidate.relative_to(root))}
    except Exception as exc:
        raise HTTPException(500, str(exc))


# ─── Query / Chat ─────────────────────────────────────────────────────────────
FILE_CTX_LIMIT    = 5_000
OVERALL_BUDGET    = 14_000
SAMPLE_FILE_LIMIT = 3


@app.post("/api/query")
async def query(req: QueryRequest):
    p1        = Path(state["phase1_output"])
    gemma_up  = True  # Cloud API always available
    has_files = bool(req.orig_file or req.trans_file)

    if req.mode != "ir":
        try:
            answer = await _smart_gemma_query(req, p1)
            return {"answer": answer, "mode": "gemma",
                    "context_type": "file" if has_files else "whole_repo"}
        except Exception as e:
            _log(f"Gemma query error -> IR fallback: {e}")

    answer = _ir_search_formatted(req.question, p1, gemma_up, has_files)
    return {"answer": answer, "mode": "ir",
            "context_type": "file" if has_files else "whole_repo"}


def _get_console_context(limit: int = 30) -> str:
    """Extracts the last N lines from the global log buffer to provide context for build failures."""
    logs = list(state.get("log_buffer", []))
    if not logs:
        return ""
    
    last_logs = logs[-limit:]
    formatted = "\n".join(last_logs)
    return f"## CURRENT CONSOLE LOGS (Last Build/Transform Output)\n```text\n{formatted}\n```"


async def _smart_gemma_query(req: QueryRequest, p1: Path) -> str:
    model  = state.get("gemma_model", "gemma-4-31b-it")
    
    parts: list[str] = []
    budget = OVERALL_BUDGET

    # 1. ALWAYS include the Architecture Map
    ir_overview = _build_ir_overview(p1)
    if ir_overview:
        parts.append(f"## PROJECT ARCHITECTURE MAP\n{ir_overview}")
        budget -= len(ir_overview)

    # 2. Add Recent Console Logs (Critical for fixing build errors)
    console_logs = _get_console_context(35)
    if console_logs:
        parts.append(console_logs)
        budget -= len(console_logs)

    # 3. Add Phase 2 Document Embeddings
    p2_context = _read_phase2_docs(p1, req.question, budget // 3)
    if p2_context:
        parts.append(p2_context)
        budget -= len(p2_context)

    # 3. Append Context Based on Mode
    if req.orig_file or req.trans_file:
        # FILE MODE
        orig_txt  = _read_file_content(req.orig_file,  "original",    6000)
        trans_txt = _read_file_content(req.trans_file, "transformed", 6000)

        if orig_txt:
            parts.append(f"### CURRENT JAVA FILE ({req.orig_file})\n```java\n{orig_txt}\n```")
            budget -= len(orig_txt)

        if trans_txt:
            lang  = "python" if (req.trans_file or "").endswith(".py") else "java"
            parts.append(f"### CURRENT TRANSFORMED FILE\n```{lang}\n{trans_txt}\n```")
            budget -= len(trans_txt)
    else:
        # WHOLE REPO MODE
        repo_files = _gather_repo_sample(req.question, p1, budget)
        if repo_files:
            parts.append(repo_files)
        elif not ir_overview:
            parts.append("## Note\nPhase 1 IR is not available. Please run Phase 1 first for richer answers.")

    parts.append(f"## USER QUESTION\n{req.question}")
    prompt = "\n\n---\n\n".join(parts)

    system = (
        "You are 'CodeModernizer AI', an expert software engineer and architect. "
        "You are assisting with modernization (Java to Java 21/Python). "
        "CRITICAL RULE: You MUST politely decline any question that is not directly "
        "related to the current codebase, Java/Python modernization, or the provided "
        "architecture map. Do not answer questions about general career advice, "
        "building a portfolio, or topics outside the immediate codebase scope.\n"
        "If no specific file is selected, answer using the 'Project Architecture Map'. "
        "Explain high-level patterns, dependencies, and entry points. "
        "Give specific, concrete answers."
    )
    
    return await _gemma_call(prompt, model, system)


def _build_ir_overview(p1: Path) -> str:
    ir_path = p1 / "ir_registry.json"
    if not ir_path.exists():
        return ""
    try:
        ir   = json.loads(ir_path.read_text())
        dep  = json.loads((p1 / "dependency_graph.json").read_text()) \
               if (p1 / "dependency_graph.json").exists() else {}
        meta = json.loads((p1 / "project_metadata.json").read_text()) \
               if (p1 / "project_metadata.json").exists() else {}

        modules  = ir.get("modules", [])
        procs    = ir.get("procedures", [])
        structs  = ir.get("data_structures", [])
        deps     = dep.get("dependencies", [])
        cycles   = meta.get("circular_dependencies", [])
        langs    = meta.get("language_distribution", {})
        entries  = set(meta.get("entry_points", []))

        lines = [
            "## Phase 1 — Project IR Overview",
            f"- {len(modules)} classes  |  {len(procs)} methods  |  "
            f"{len(structs)} data structures  |  {len(deps)} dependencies",
            f"- Languages: {langs}",
        ]

        if entries:
            entry_names = [m["name"] for m in modules if m["module_id"] in entries]
            lines.append(f"- Entry points: {', '.join(entry_names)}")

        if cycles:
            lines.append(f"- Circular deps: {len(cycles)} cycle(s) detected")

        lines.append("")
        lines.append("### All Classes (name → file)")
        for m in sorted(modules, key=lambda x: x.get("name", "")):
            ep    = " ★" if m.get("is_entry_point") else ""
            lines.append(f"  {m['name']}{ep}  ←  {m.get('file_path','?')}")

        return "\n".join(lines)
    except Exception:
        return ""


def _read_phase2_docs(p1: Path, question: str, budget: int) -> str:
    p2_storage = Path(state["phase2_storage"])
    docstore   = p2_storage / "docstore.json"
    if not docstore.exists():
        return ""

    try:
        raw      = json.loads(docstore.read_text())
        docs_data = raw.get("docstore/data", raw)
        texts: list[tuple] = []
        q_words = set(question.lower().split())

        for _id, node in docs_data.items():
            text = ""
            if isinstance(node, dict):
                data = node.get("__data__", node)
                text = data.get("text", "") or data.get("content", "")
            if not text:
                continue
            text_lower = text.lower()
            score = sum(1 for w in q_words if len(w) > 3 and w in text_lower)
            texts.append((score, text))

        texts.sort(key=lambda x: -x[0])
        selected: list[str] = []
        used = 0
        for score, text in texts[:20]:
            chunk = text[:1500]
            if used + len(chunk) > budget:
                break
            selected.append(chunk)
            used += len(chunk)

        if not selected:
            return ""

        return "## Phase 2 — Architectural Summaries\n" + "\n---\n".join(selected)
    except Exception:
        return ""


def _gather_repo_sample(question: str, p1: Path, budget: int) -> str:
    """
    For whole-repo queries: read entry-point files + keyword-matched files.
    Works even when repo_path is None (uses transformed output as fallback).
    """
    ir_path = p1 / "ir_registry.json"
    if not ir_path.exists():
        return ""

    repo_root  = Path(state["repo_path"]) if state["repo_path"] else None
    out_root   = Path(state["transform_output"])
    # If original source is gone, try the transform output
    if not repo_root or not repo_root.exists():
        repo_root = out_root if out_root.exists() else None

    try:
        ir      = json.loads(ir_path.read_text())
        meta    = json.loads((p1 / "project_metadata.json").read_text()) \
                  if (p1 / "project_metadata.json").exists() else {}
        modules = ir.get("modules", [])
        entries = set(meta.get("entry_points", []))
        words   = [w.lower() for w in question.split() if len(w) > 3]

        priority: list[dict] = []
        others:   list[dict] = []
        for m in modules:
            if m["module_id"] in entries:
                priority.append(m)
            elif words and any(w in m.get("name", "").lower() for w in words):
                priority.append(m)
            else:
                others.append(m)

        candidates = (priority[:SAMPLE_FILE_LIMIT] +
                      others[:max(0, SAMPLE_FILE_LIMIT - len(priority))])
        blocks: list[str] = []
        used = 0

        for m in candidates[:SAMPLE_FILE_LIMIT]:
            fp = m.get("file_path", "")
            if not fp:
                continue

            # Try to read from original repo
            if repo_root:
                orig_path = repo_root / fp
                if orig_path.exists():
                    txt = orig_path.read_text(encoding="utf-8", errors="replace")[:FILE_CTX_LIMIT]
                    block = f"### Original: `{fp}`\n```java\n{txt}\n```"
                    if used + len(block) < budget:
                        blocks.append(block)
                        used += len(block)

            # Try transformed counterpart
            if out_root and out_root.exists():
                stem = Path(fp).stem
                for ext in (".py", ".java"):
                    matches = list(out_root.rglob(f"{stem}{ext}"))
                    if matches:
                        txt2 = matches[0].read_text(encoding="utf-8", errors="replace")[:FILE_CTX_LIMIT]
                        lang = "python" if ext == ".py" else "java"
                        block2 = f"### Transformed: `{matches[0].name}`\n```{lang}\n{txt2}\n```"
                        if used + len(block2) < budget:
                            blocks.append(block2)
                            used += len(block2)
                        break

        if not blocks:
            return ""

        return "## Sample Source Files (original + transformed)\n\n" + "\n\n".join(blocks)
    except Exception:
        return ""


# ─── IR search (no LLM) ──────────────────────────────────────────────────────
def _read_file_content(rel_path: Optional[str], kind: str, limit: int) -> str:
    if not rel_path:
        return ""
    root = (Path(state["repo_path"]) if kind == "original" and state["repo_path"]
            else Path(state["transform_output"]))
    if not root or not root.exists():
        return ""
    candidate = root / rel_path
    if not candidate.exists():
        swapped = candidate.with_suffix(".py" if candidate.suffix == ".java" else ".java")
        candidate = swapped if swapped.exists() else None
    if not candidate or not candidate.exists():
        return ""
    try:
        txt = candidate.read_text(encoding="utf-8", errors="replace")
        return txt[:limit] + ("\n# [truncated]" if len(txt) > limit else "")
    except Exception:
        return ""


def _ir_search_formatted(question: str, p1: Path, gemma_up: bool, has_files: bool) -> str:
    ir_path = p1 / "ir_registry.json"
    if not ir_path.exists():
        return (
            "⚠️ **Phase 1 has not been run yet.**\n\n"
            "Run Phase 1 first to enable code search and AI-powered answers."
        )

    ir   = json.loads(ir_path.read_text())
    dep  = json.loads((p1 / "dependency_graph.json").read_text()) \
           if (p1 / "dependency_graph.json").exists() else {}
    meta = json.loads((p1 / "project_metadata.json").read_text()) \
           if (p1 / "project_metadata.json").exists() else {}

    words   = [w.lower() for w in question.split() if len(w) > 2]
    results: list[str] = []

    for m in ir.get("modules", []):
        if not words or any(w in m.get("name", "").lower() for w in words):
            ep = " ★ entry-point" if m.get("is_entry_point") else ""
            results.append(f"📦 **{m['name']}** ({m.get('language','?')})  `{m.get('file_path','')}`{ep}")

    for p in ir.get("procedures", []):
        if words and any(w in p.get("name", "").lower() for w in words):
            params = ", ".join(p.get("parameters", []))
            results.append(f"⚙️ **{p['name']}**({params}) → `{p.get('return_type') or 'void'}`  `{p.get('source_file','?')}`")

    modules  = ir.get("modules", [])
    procs    = ir.get("procedures", [])
    structs  = ir.get("data_structures", [])
    dep_list = dep.get("dependencies", [])
    cycles   = meta.get("circular_dependencies", [])

    gemma_note = (
        ""
        if gemma_up else "\n\n⚡ **Gemma API unavailable** — check your API key for full AI-powered answers."
    )

    OVERVIEW_WORDS = {"what","does","this","repo","project","about","explain","describe",
                      "overview","purpose","summary","whole","codebase","entire","all",
                      "system","application","app","how","work","works","do","doing"}
    question_words_set = set(question.lower().split())
    is_overview = (
        not words
        or len(question_words_set - OVERVIEW_WORDS) <= 1
        or any(w in question.lower() for w in ["what does", "what is this", "what do",
                                                "describe", "overview", "explain"])
    )

    if is_overview:
        return _ir_project_summary(ir, dep, meta, modules, procs, structs,
                                   dep_list, cycles, gemma_note)

    header = (
        f"**IR Search** — {len(modules)} classes · {len(procs)} methods · "
        f"{len(structs)} structs · {len(dep_list)} deps\n\n"
    )
    if results:
        return header + "\n".join(results[:30]) + gemma_note

    return _ir_project_summary(ir, dep, meta, modules, procs, structs,
                               dep_list, cycles, gemma_note)


def _ir_project_summary(ir, dep, meta, modules, procs, structs,
                        dep_list, cycles, gemma_note="") -> str:
    langs   = meta.get("language_distribution", {})
    entries = set(meta.get("entry_points", []))
    entry_modules = [m for m in modules if m["module_id"] in entries]
    entry_names   = [m["name"] for m in entry_modules]

    DOMAINS: dict[str, list[str]] = {}
    domain_keywords = {
        "Payment / Billing": ["payment", "billing", "invoice", "currency", "price", "money", "charge"],
        "Customer / User":   ["customer", "user", "account", "client", "profile", "person"],
        "Reporting":         ["report", "analytics", "audit", "log", "export", "generate"],
        "Data / Storage":    ["database", "pool", "repository", "store", "cache", "persist", "dao"],
        "Scheduling":        ["scheduler", "timer", "cron", "task", "job", "queue"],
        "Notification":      ["notification", "email", "message", "alert", "notify"],
        "Product / Catalog": ["product", "catalog", "item", "sku", "inventory"],
        "Security / Fraud":  ["fraud", "security", "auth", "detector", "checker"],
        "Core / Utility":    [],
    }
    assigned: set[str] = set()
    for domain, keywords in domain_keywords.items():
        if not keywords:
            continue
        matched = [m["name"] for m in modules
                   if any(k in m.get("name", "").lower() for k in keywords)]
        if matched:
            DOMAINS[domain] = matched
            assigned.update(matched)

    unassigned = [m["name"] for m in modules if m["name"] not in assigned]
    if unassigned:
        DOMAINS["Core / Utility"] = unassigned

    dep_counts: dict[str, int] = {}
    for d in dep_list:
        src = d.get("source_module", "")
        dep_counts[src] = dep_counts.get(src, 0) + 1
    most_depended = sorted(dep_counts.items(), key=lambda x: -x[1])[:5]

    p3 = Path(state["transform_output"])
    has_python = p3.exists() and any(p3.rglob("*.py"))
    has_java21 = p3.exists() and any(p3.rglob("*.java")) and not has_python
    transform_note = ""
    if has_python:
        transform_note = f"\n\n**Transformation:** ✓ Translated to Python 3.11+ ({len(list(p3.rglob('*.py')))} files in `{p3.name}/`)"
    elif has_java21:
        transform_note = f"\n\n**Transformation:** ✓ Modernised to Java 21 ({len(list(p3.rglob('*.java')))} files in `{p3.name}/`)"

    lines: list[str] = []
    lines.append(f"**Project Overview** — {len(modules)} classes · {len(procs)} methods · "
                 f"{len(structs)} data structures · {len(dep_list)} dependencies")

    if langs:
        lang_str = ", ".join(f"{k} ({v})" for k, v in langs.items())
        lines.append(f"**Languages:** {lang_str}")

    if entry_names:
        lines.append(f"**Entry points:** {', '.join(entry_names)}")

    if cycles:
        lines.append(f"**⚠ Circular dependencies:** {len(cycles)} cycle(s) detected")

    lines.append("")
    lines.append("**Functional Areas:**")
    for domain, names in DOMAINS.items():
        if names:
            lines.append(f"  • **{domain}**: {', '.join(sorted(names))}")

    if most_depended:
        lines.append("")
        lines.append("**Most connected classes (dependency hubs):**")
        for name, count in most_depended:
            lines.append(f"  • {name} — {count} outgoing deps")

    lines.append("")
    purpose_parts: list[str] = []
    if DOMAINS.get("Payment / Billing"):   purpose_parts.append("billing and payment processing")
    if DOMAINS.get("Customer / User"):     purpose_parts.append("customer management")
    if DOMAINS.get("Reporting"):           purpose_parts.append("reporting and auditing")
    if DOMAINS.get("Data / Storage"):      purpose_parts.append("data persistence")
    if DOMAINS.get("Scheduling"):          purpose_parts.append("scheduled task execution")
    if DOMAINS.get("Notification"):        purpose_parts.append("notification delivery")
    if DOMAINS.get("Product / Catalog"):   purpose_parts.append("product/catalog management")
    if DOMAINS.get("Security / Fraud"):    purpose_parts.append("fraud detection and security")

    if purpose_parts:
        purpose = ", ".join(purpose_parts[:-1])
        if len(purpose_parts) > 1:
            purpose += f", and {purpose_parts[-1]}"
        else:
            purpose = purpose_parts[0]
        lines.append(f"**Inferred purpose:** This appears to be a **{purpose}** system.")
    else:
        lines.append(f"**{len(modules)} classes** across {len(set(DOMAINS.keys()))} functional areas.")

    lines.append(transform_note)
    lines.append(gemma_note)

    return "\n".join(l for l in lines if l is not None)


# ─── Frontend ─────────────────────────────────────────────────────────────────
@app.get("/")
async def serve_ui():
    html_file = PROJECT_ROOT / "frontend" / "index.html"
    if not html_file.exists():
        return HTMLResponse("<h2>frontend/index.html not found</h2>")
    return HTMLResponse(html_file.read_text(encoding="utf-8"))


if __name__ == "__main__":
    PORT = int(os.environ.get("PORT", 8000))
    try:
        from pyngrok import ngrok, conf
        token = os.environ.get("NGROK_TOKEN", "")
        if token:
            conf.get_default().auth_token = token
        tunnel = ngrok.connect(PORT, "http")
        print(f"\n{'='*55}\n  🌐  Public URL  : {tunnel.public_url}\n  📍  Local  URL  : http://localhost:{PORT}\n{'='*55}\n")
    except Exception:
        print(f"\n  📍  Running at: http://localhost:{PORT}\n")
    uvicorn.run(app, host="0.0.0.0", port=PORT, log_level="warning")