pipeline.py

import json
import re
import asyncio
import traceback

from agents import ResearchAgent, OrchestratorAgent, FrontendAgent, BackendAgent
from schemas import ProjectState, AgentMessage


class PipelineEngine:
    def __init__(self):
        self.research = ResearchAgent()
        self.orchestrator = OrchestratorAgent()
        self.frontend = FrontendAgent()
        self.backend = BackendAgent()

    # ── helpers ───────────────────────────────────────────
    def _emit(
        self,
        state: ProjectState,
        event_type: str,
        agent: str,
        content: str = "",
        file_path: str | None = None,
        metadata: dict | None = None,
    ):
        state.messages.append(
            AgentMessage(
                event_type=event_type,
                agent=agent,
                content=content,
                file_path=file_path,
                metadata=metadata or {},
            )
        )

    # ── main pipeline ────────────────────────────────────
    async def run(self, state: ProjectState):
        try:
            # ── 1. Research ──────────────────────────────
            state.status = "researching"
            state.current_agent = "research"
            self._emit(state, "agent_start", "research", "🔍 Starting research…")

            research = await self._do_research(state)
            state.research_output = research
            self._emit(
                state,
                "agent_done",
                "research",
                "✅ Research complete",
                metadata={"keys": list(research.keys())},
            )

            # ── 2. Orchestrate ───────────────────────────
            state.status = "orchestrating"
            state.current_agent = "orchestrator"
            self._emit(state, "agent_start", "orchestrator", "🧠 Building blueprint…")

            blueprint = await self._do_orchestrate(state)
            state.blueprint = blueprint
            self._emit(
                state,
                "agent_done",
                "orchestrator",
                "✅ Blueprint ready",
                metadata={"systems": list(blueprint.get("systems", {}).keys())},
            )

            # ── 3 & 4. Frontend + Backend in parallel ────
            state.status = "building"
            state.current_agent = "frontend+backend"
            self._emit(state, "agent_start", "system", "🚀 Generating code…")

            await asyncio.gather(
                self._do_frontend(state),
                self._do_backend(state),
            )

            # ── 5. Merge ────────────────────────────────
            state.status = "merging"
            state.current_agent = "system"
            self._emit(state, "agent_start", "system", "📦 Merging…")

            state.file_tree = sorted(state.generated_files.keys())
            self._create_preview(state)

            state.status = "completed"
            state.current_agent = ""
            self._emit(
                state,
                "agent_done",
                "system",
                f"🎉 Done — {len(state.generated_files)} files generated",
            )

        except Exception as exc:
            state.status = "error"
            state.errors.append(str(exc))
            self._emit(
                state,
                "error",
                state.current_agent or "system",
                f"❌ {exc}",
                metadata={"tb": traceback.format_exc()},
            )

    # ── research phase ───────────────────────────────────
    async def _do_research(self, state: ProjectState) -> dict:
        prompt = (
            f"Research this app idea and produce a JSON report:\n\n"
            f"APP IDEA: {state.user_prompt}\n"
            f"APP TYPE: {state.app_type}\n"
            f"SYSTEMS: {', '.join(state.systems)}\n\n"
            "Return JSON with keys: stack_recommendation, schema_hints, "
            "api_docs_summary, security_notes, hosting_notes, ui_patterns, "
            "competitor_analysis"
        )
        buf: list[str] = []
        async for tok in self.research.call([{"role": "user", "content": prompt}]):
            buf.append(tok)
            if len(buf) % 30 == 0:
                self._emit(state, "token", "research", tok)
        text = "".join(buf)
        self._emit(state, "token", "research", "\n[done]")
        return self._parse_json(text)

    # ── orchestrate phase ────────────────────────────────
    async def _do_orchestrate(self, state: ProjectState) -> dict:
        prompt = (
            f"Create a master blueprint for:\n\n"
            f"USER REQUEST: {state.user_prompt}\n"
            f"APP TYPE: {state.app_type}\n"
            f"SYSTEMS: {', '.join(state.systems)}\n\n"
            f"RESEARCH:\n{json.dumps(state.research_output, indent=2, default=str)[:6000]}\n\n"
            "Produce a complete master blueprint JSON following your instructions."
        )
        buf: list[str] = []
        async for tok in self.orchestrator.call([{"role": "user", "content": prompt}]):
            buf.append(tok)
            if len(buf) % 20 == 0:
                self._emit(state, "token", "orchestrator", tok)
        text = "".join(buf)
        return self._parse_json(text)

    # ── frontend phase ───────────────────────────────────
    async def _do_frontend(self, state: ProjectState):
        self._emit(state, "agent_start", "frontend", "🎨 Generating frontend…")

        for system in state.systems:
            sys_bp = state.blueprint.get("systems", {}).get(system, {})
            design = state.blueprint.get("design_tokens", {})
            pay = state.blueprint.get("payment_config", {})

            prompt = (
                f'Generate complete frontend for the "{system}" system.\n\n'
                f"BLUEPRINT:\n{json.dumps(sys_bp, indent=2, default=str)}\n\n"
                f"DESIGN TOKENS:\n{json.dumps(design, indent=2, default=str)}\n\n"
                f"PAYMENT CONFIG:\n{json.dumps(pay, indent=2, default=str)}\n\n"
                f"DB SCHEMA (ref):\n"
                f"{json.dumps(state.blueprint.get('database_schema',{}), indent=2, default=str)[:3000]}\n\n"
                f"Mark every file: // FILE: {system}/path/file.jsx\n"
                "Generate ALL files. Use React + Tailwind."
            )
            buf: list[str] = []
            async for tok in self.frontend.call([{"role": "user", "content": prompt}]):
                buf.append(tok)
                if len(buf) % 15 == 0:
                    self._emit(state, "token", "frontend", tok)
            text = "".join(buf)
            for path, content in self._extract_files(text).items():
                state.generated_files[path] = content
                self._emit(state, "file_created", "frontend", f"📄 {path}", file_path=path)

        self._emit(state, "agent_done", "frontend", "✅ Frontend complete")

    # ── backend phase ────────────────────────────────────
    async def _do_backend(self, state: ProjectState):
        self._emit(state, "agent_start", "backend", "🔐 Generating backend…")

        prompt = (
            f"Generate complete backend for this app.\n\n"
            f"DB SCHEMA:\n{json.dumps(state.blueprint.get('database_schema',{}), indent=2, default=str)}\n\n"
            f"API ENDPOINTS:\n{json.dumps(state.blueprint.get('api_endpoints',[]), indent=2, default=str)}\n\n"
            f"AUTH CONFIG:\n{json.dumps(state.blueprint.get('auth_config',{}), indent=2, default=str)}\n\n"
            f"PAYMENT CONFIG:\n{json.dumps(state.blueprint.get('payment_config',{}), indent=2, default=str)}\n\n"
            f"APP: {state.user_prompt}\n\n"
            "Generate: SQL schema, FastAPI backend, Edge Functions, Docker files, "
            ".env.example, firebase.json, DEPLOY.md.\n"
            "Mark files: # FILE: path  or  -- FILE: path"
        )
        buf: list[str] = []
        async for tok in self.backend.call([{"role": "user", "content": prompt}]):
            buf.append(tok)
            if len(buf) % 15 == 0:
                self._emit(state, "token", "backend", tok)
        text = "".join(buf)
        for path, content in self._extract_files(text).items():
            state.generated_files[path] = content
            self._emit(state, "file_created", "backend", f"📄 {path}", file_path=path)

        self._emit(state, "agent_done", "backend", "✅ Backend complete")

    # ── bug fix ──────────────────────────────────────────
    async def fix(self, state: ProjectState, error_msg: str, file_path: str | None):
        self._emit(state, "agent_start", "backend", f"🔧 Fixing: {error_msg[:80]}")

        code = ""
        if file_path and file_path in state.generated_files:
            code = state.generated_files[file_path][:8000]

        prompt = (
            f"FIX THIS BUG:\nERROR: {error_msg}\n"
            + (f"FILE: {file_path}\nCODE:\n{code}\n" if code else "")
            + "\nOutput the COMPLETE fixed file with original file path marker."
        )

        is_fe = file_path and any(
            file_path.endswith(e) for e in (".jsx", ".tsx", ".css", ".html", ".js")
        )
        agent = self.frontend if is_fe else self.backend
        name = "frontend" if is_fe else "backend"

        buf: list[str] = []
        async for tok in agent.call([{"role": "user", "content": prompt}]):
            buf.append(tok)
        text = "".join(buf)
        for p, c in self._extract_files(text).items():
            state.generated_files[p] = c
            self._emit(state, "file_created", name, f"🔧 Fixed {p}", file_path=p)

        state.status = "completed"
        self._emit(state, "agent_done", name, "✅ Fix applied")

    # ── JSON parser ──────────────────────────────────────
    def _parse_json(self, text: str) -> dict:
        for pattern in [
            r"```json\s*([\s\S]*?)\s*```",
            r"```\s*([\s\S]*?)\s*```",
        ]:
            for m in re.findall(pattern, text):
                try:
                    return json.loads(m)
                except json.JSONDecodeError:
                    continue
        # try raw
        brace = text.find("{")
        if brace != -1:
            depth, end = 0, brace
            for i in range(brace, len(text)):
                if text[i] == "{":
                    depth += 1
                elif text[i] == "}":
                    depth -= 1
                    if depth == 0:
                        end = i + 1
                        break
            try:
                return json.loads(text[brace:end])
            except json.JSONDecodeError:
                pass

        return self._fallback_blueprint()

    def _fallback_blueprint(self) -> dict:
        return {
            "project_name": "generated-app",
            "description": "AI-generated web application",
            "systems": {
                "client_portal": {
                    "pages": [
                        {"path": "/dashboard", "title": "Dashboard", "components": ["Layout", "Stats", "Activity"], "auth_required": True},
                        {"path": "/settings", "title": "Settings", "components": ["ProfileForm", "PasswordForm"], "auth_required": True},
                        {"path": "/billing", "title": "Billing", "components": ["PlanCards", "PayPalButton", "InvoiceList"], "auth_required": True},
                    ],
                    "features": ["auth", "dashboard", "billing", "settings"],
                },
                "public_landing": {
                    "pages": [
                        {"path": "/", "title": "Home", "components": ["Hero", "Features", "Pricing", "CTA", "Footer"], "auth_required": False},
                        {"path": "/pricing", "title": "Pricing", "components": ["PricingCards", "FAQ"], "auth_required": False},
                    ],
                    "features": ["hero", "pricing", "signup", "seo"],
                },
                "marketing_cms": {
                    "pages": [
                        {"path": "/blog", "title": "Blog", "components": ["BlogList", "PostEditor"], "auth_required": True},
                    ],
                    "features": ["blog", "email_capture"],
                },
                "analytics_dashboard": {
                    "pages": [
                        {"path": "/analytics", "title": "Analytics", "components": ["Charts", "Metrics", "UserTable"], "auth_required": True},
                    ],
                    "features": ["realtime_metrics", "charts", "revenue"],
                },
                "admin_panel": {
                    "pages": [
                        {"path": "/admin", "title": "Admin", "components": ["UserMgmt", "Logs", "SystemHealth"], "auth_required": True},
                    ],
                    "features": ["user_management", "moderation", "logs"],
                },
            },
            "database_schema": {"tables": []},
            "api_endpoints": [],
            "auth_config": {"providers": ["email"]},
            "payment_config": {
                "provider": "paypal",
                "plans": [
                    {"name": "Free", "price": 0, "features": ["Basic"]},
                    {"name": "Pro", "price": 29, "features": ["All features"]},
                ],
            },
            "design_tokens": {
                "colors": {
                    "primary": "#6C63FF",
                    "secondary": "#00D9FF",
                    "background": "#0A0A0F",
                    "surface": "#111118",
                    "text": "#F0F0FF",
                },
                "fonts": {"heading": "Inter", "body": "Inter", "mono": "JetBrains Mono"},
            },
        }

    # ── file extractor ───────────────────────────────────
    def _extract_files(self, text: str) -> dict[str, str]:
        files: dict[str, str] = {}
        pattern = r"(?://|#|--)\s*FILE:\s*(.+?)(?:\n)([\s\S]*?)(?=(?://|#|--)\s*FILE:|$)"
        for path, content in re.findall(pattern, text):
            clean = path.strip()
            body = re.sub(r"\s*```\s*$", "", content.strip())
            body = re.sub(r"^```\w*\s*", "", body.strip())
            if clean and body:
                files[clean] = body

        if not files:
            blocks = re.findall(
                r"(?:#+\s*)?(?:`([^`]+)`|(\S+\.\w+))\s*\n```\w*\n([\s\S]*?)```", text
            )
            for n1, n2, body in blocks:
                name = (n1 or n2 or "").strip()
                if name and body.strip():
                    files[name] = body.strip()

        if not files:
            if "def " in text or "import " in text:
                files["backend/generated.py"] = text
            elif "CREATE TABLE" in text.upper():
                files["database/schema.sql"] = text
            elif "function" in text or "const " in text:
                files["frontend/generated.jsx"] = text
            else:
                files["output/raw.txt"] = text

        return files

    # ── preview builder ──────────────────────────────────
    def _create_preview(self, state: ProjectState):
        bp = state.blueprint
        colors = bp.get("design_tokens", {}).get("colors", {})
        name = bp.get("project_name", "Generated App")
        systems = bp.get("systems", {})

        cards = ""
        for sname, sdata in systems.items():
            pages = sdata.get("pages", [])
            feats = sdata.get("features", [])
            nice = sname.replace("_", " ").title()
            feat_html = "".join(f'<span class="t">{f}</span>' for f in feats[:6])
            page_html = "".join(
                f'<li>{p.get("title", p.get("path", ""))}</li>' for p in pages[:6]
            )
            cards += (
                f'<div class="c"><h3>{nice}</h3>'
                f'<div class="ts">{feat_html}</div>'
                f"<ul>{page_html}</ul>"
                f'<small>{len(pages)} pages</small></div>'
            )

        fc = len(state.generated_files)
        ext_counts: dict[str, int] = {}
        for f in state.generated_files:
            e = f.rsplit(".", 1)[-1] if "." in f else "other"
            ext_counts[e] = ext_counts.get(e, 0) + 1
        fstats = " · ".join(f"{v} .{k}" for k, v in sorted(ext_counts.items()))
        total_pages = sum(len(s.get("pages", [])) for s in systems.values())

        pri = colors.get("primary", "#6C63FF")
        sec = colors.get("secondary", "#00D9FF")
        bg = colors.get("background", "#0A0A0F")
        sf = colors.get("surface", "#111118")
        tx = colors.get("text", "#F0F0FF")

        html = f"""<!DOCTYPE html>
<html lang="en"><head><meta charset="UTF-8"><meta name="viewport" content="width=device-width,initial-scale=1">
<title>{name}</title>
<link href="https://fonts.googleapis.com/css2?family=Inter:wght@300;400;600;700&family=JetBrains+Mono&display=swap" rel="stylesheet">
<style>
*{{margin:0;padding:0;box-sizing:border-box}}
body{{font-family:'Inter',sans-serif;background:{bg};color:{tx};min-height:100vh;padding:2rem}}
.w{{max-width:1100px;margin:0 auto}}
.hero{{text-align:center;padding:3.5rem 2rem;background:linear-gradient(135deg,{sf},{bg});border-radius:20px;border:1px solid #1E1E2E;margin-bottom:2rem}}
.hero h1{{font-size:2.6rem;font-weight:700;background:linear-gradient(135deg,{pri},{sec});-webkit-background-clip:text;-webkit-text-fill-color:transparent;margin-bottom:.8rem}}
.hero p{{color:#8888AA;font-size:1.1rem;max-width:550px;margin:0 auto}}
.stats{{display:flex;gap:1rem;justify-content:center;margin-top:1.5rem;flex-wrap:wrap}}
.s{{background:{sf};border:1px solid #1E1E2E;border-radius:12px;padding:.8rem 1.3rem;text-align:center}}
.sv{{font-size:1.8rem;font-weight:700;color:{pri}}}
.sl{{color:#8888AA;font-size:.8rem;margin-top:.15rem}}
.g{{display:grid;grid-template-columns:repeat(auto-fit,minmax(280px,1fr));gap:1.2rem;margin-top:1.5rem}}
.c{{background:{sf};border:1px solid #1E1E2E;border-radius:14px;padding:1.3rem;transition:.3s}}
.c:hover{{border-color:{pri}44;transform:translateY(-2px);box-shadow:0 6px 24px {pri}15}}
.c h3{{font-size:1.1rem;font-weight:600;margin-bottom:.6rem}}
.ts{{display:flex;flex-wrap:wrap;gap:.3rem;margin-bottom:.8rem}}
.t{{background:{pri}22;color:{pri};padding:.15rem .5rem;border-radius:5px;font-size:.7rem;font-weight:500}}
ul{{list-style:none;margin-bottom:.5rem}}
li{{padding:.25rem 0;color:#8888AA;font-size:.85rem;border-bottom:1px solid #1E1E2E}}
li:last-child{{border:none}}
small{{color:{sec};font-size:.75rem}}
.fi{{text-align:center;margin-top:1.5rem;color:#8888AA;font-family:'JetBrains Mono',monospace;font-size:.8rem}}
.sec{{font-size:1.3rem;font-weight:600;margin-bottom:.4rem}}
@keyframes fadeIn{{from{{opacity:0;transform:translateY(8px)}}to{{opacity:1;transform:translateY(0)}}}}
.c{{animation:fadeIn .5s ease forwards}}
.c:nth-child(2){{animation-delay:.1s}}.c:nth-child(3){{animation-delay:.2s}}
.c:nth-child(4){{animation-delay:.3s}}.c:nth-child(5){{animation-delay:.4s}}
</style></head>
<body><div class="w">
<div class="hero"><h1>{name}</h1><p>{state.user_prompt[:180]}</p>
<div class="stats">
<div class="s"><div class="sv">{len(systems)}</div><div class="sl">Systems</div></div>
<div class="s"><div class="sv">{fc}</div><div class="sl">Files</div></div>
<div class="s"><div class="sv">{total_pages}</div><div class="sl">Pages</div></div>
</div></div>
<h2 class="sec">Generated Systems</h2>
<div class="g">{cards}</div>
<div class="fi">{fstats}</div>
</div></body></html>"""

        state.generated_files["preview/index.html"] = html