"""
NeuraPrompt Studio — autonomous web-app builder (Lovable-style).

Key design:
  1. PLANNER pass  -> AI returns a strict JSON manifest of files (path + purpose
     + dependencies). It decides the stack, architecture, design and feature
     set on its own, given only the user's prompt.
  2. PER-FILE pass -> Each file in the manifest is generated strictly sequentially 
     to ensure free-tier models don't trigger concurrent abuse blocks or rate limits.
  3. CONTEXT INJECTION -> Each subsequent file generation receives the complete raw 
     source code of all previously written files to eliminate broken connections 
     (e.g., missing CSS classes, unhooked hamburger menus, static dynamic hooks).
  4. CONTINUATION -> If a single file response hits the model's max_tokens
     (finish_reason == "length"), we automatically issue a "continue exactly
     where you left off, no preamble" call and stitch the chunks together.
  5. WIRING -> After all files are written, vanilla HTML projects auto-link
     CSS / JS so the preview works without manual edits.
"""

import json
import os
import re
import time
import asyncio
import requests
from datetime import datetime, timedelta
from collections import defaultdict, deque

from fastapi import APIRouter, Request
from fastapi.responses import StreamingResponse, JSONResponse
from motor.motor_asyncio import AsyncIOMotorClient
from fastapi.middleware.cors import CORSMiddleware
from bson import ObjectId

# FastAPI
from fastapi import FastAPI, Form, HTTPException, Query, UploadFile, File, Request
from fastapi.responses import StreamingResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field


router = APIRouter(prefix="/api/v1/studio", tags=["studio"])

# ── ENV ───────────────────────────────────────────────────────────────────────
OPENROUTER_KEY = os.environ.get("OPENROUTE_KEY")
MONGO_URL      = os.environ.get("MONGO_URI")

mongo_client = AsyncIOMotorClient(MONGO_URL)
db           = mongo_client["neuraprompt"]

# ── MODELS ────────────────────────────────────────────────────────────────────
BUILD_MODEL    = "openrouter/owl-alpha"
FAST_MODEL     = "poolside/laguna-xs.2:free"
FALLBACK_MODEL = "poolside/laguna-m.1:free"

MAX_TOKENS_PER_CALL  = 18000
RATE_LIMIT_WAIT      = 45  # Balanced buffer window for free tier API cooling
MAX_BUILD_RETRIES    = 3
MAX_CONTINUATIONS    = 6     # per file
MAX_FILES_PER_BUILD  = 25    # safety ceiling
API_TIMEOUT_WINDOW   = 45    # Reduced from 180 to prevent multi-minute freezes on hanging endpoints

# ── TTL INDEX ─────────────────────────────────────────────────────────────────
@router.on_event("startup")
async def startup():
    try:
        await db.studio_projects.create_index("expires_at", expireAfterSeconds=0)
        await db.studio_files.create_index("project_id")
        await db.studio_files.create_index("expires_at", expireAfterSeconds=0)
        print("[Studio] TTL + file indexes ready")
    except Exception as e:
        print(f"[Studio] Index warning: {e}")

# ── RATE LIMITER ──────────────────────────────────────────────────────────────
class RateLimiter:
    def __init__(self, max_calls: int, period: float):
        self.max_calls = max_calls
        self.period    = period
        self.calls: deque = deque()

    def is_allowed(self):
        now = time.time()
        while self.calls and self.calls[0] < now - self.period:
            self.calls.popleft()
        if len(self.calls) < self.max_calls:
            self.calls.append(now)
            return True, 0.0
        return False, round(self.period - (now - self.calls[0]), 1)

global_limiter = RateLimiter(12, 60)
user_limiters  = defaultdict(lambda: RateLimiter(3, 60))

def check_rate_limits(user_id: str):
    ok, wait = global_limiter.is_allowed()
    if not ok:
        return False, f"Studio is busy right now. Please try again in {int(wait)+1} seconds."
    ok, wait = user_limiters[user_id].is_allowed()
    if not ok:
        return False, f"You are building too fast. Please wait {int(wait)+1} seconds."
    return True, ""

# ── UTILS ─────────────────────────────────────────────────────────────────────
def clean_thinking(text: str) -> str:
    text = re.sub(r'<think>[\s\S]*?</think>',       '', text, flags=re.IGNORECASE)
    text = re.sub(r'<thinking>[\s\S]*?</thinking>', '', text, flags=re.IGNORECASE)
    text = re.sub(r'^/think\s*', '', text.strip())
    return text.strip()

def sse(type_: str, payload: dict) -> str:
    return f"data: {json.dumps({'type': type_, 'payload': payload})}\n\n"

def guess_file_type(path: str) -> str:
    ext = os.path.splitext(path)[1].lower()
    return {
        ".html": "html", ".css": "css", ".js": "javascript",
        ".ts": "typescript", ".tsx": "tsx", ".jsx": "jsx",
        ".json": "json", ".md": "markdown", ".svg": "svg",
        ".env": "env", ".toml": "toml", ".yml": "yaml", ".yaml": "yaml",
    }.get(ext, "text")

def strip_code_fences(text: str) -> str:
    text = text.strip()
    m = re.match(r'^```[\w+\-]*\n?([\s\S]*?)\n?```$', text)
    if m:
        return m.group(1).strip()
    text = re.sub(r'^```[\w+\-]*\n', '', text)
    text = re.sub(r'\n```\s*$', '', text)
    return text.strip()

# ── NON-WEB GUARD ─────────────────────────────────────────────────────────────
NON_WEB_KEYWORDS = [
    "android", "kotlin", "apk", "play store", "google play",
    "fastapi", "flask", "django", "cli ", "scraper",
    "bash script", "shell script",
]

def is_non_web_request(prompt: str) -> bool:
    p = " " + prompt.lower() + " "
    return any(kw in p for kw in NON_WEB_KEYWORDS)

# ── SYSTEM PROMPTS ────────────────────────────────────────────────────────────
PLANNER_SYSTEM = """You are Neurones Pro created by Alysium Corp Studios ZA inc by Toxic Dee Modder — You are a senior product engineer who plans complete, production-grade web applications BEFORE writing any code.

You are autonomous. You decide the stack, architecture, screens, features, copy, palette, typography and motion. The user gives you an idea — you design the product.

OUTPUT FORMAT — RETURN ONLY VALID JSON (no markdown, no preamble, no commentary):
{
  "app_name": "...",
  "summary": "1-2 sentence description of what you are building",
  "stack": "html" | "vanilla" | "react-vite",
  "design": {
    "theme": "dark" | "light",
    "palette": ["#hex", "#hex", "#hex", "#hex"],
    "font_pair": "e.g. 'Space Grotesk + Inter'",
    "vibe": "1 line — e.g. 'editorial brutalist with neon accents'"
  },
  "files": [
    {
      "path": "index.html",
      "purpose": "what this file does",
      "depends_on": ["src/styles.css", "src/main.js", "styles.css", "script.js"],
      "key_features": ["specific feature 1", "specific feature 2"]
    }
  ]
}

PLANNING RULES:
- Build a REAL product, not a demo. Multiple sections, real interactivity, real content.
- Decide the stack yourself: vanilla HTML/CSS/JS for landing pages, marketing sites, simple tools, games. React+Vite+TS for dashboards, multi-view apps, data-heavy tools, HTML for portfolio, dashboards, landing pages, simple websites, in each prompt look for html or react+Vite or vanilla, if not stated choose yourself.
- For vanilla projects: include index.html, src/styles.css, src/main.js, README.md plus any extra modules you need (components, data, utils).
- For react-vite projects: include package.json, vite.config.ts, tsconfig.json, index.html, src/main.tsx, src/App.tsx, src/index.css, plus components, hooks, types, README.md.
- For html projects: include index.html, styles.css, script.js, plus README.md.
- Minimum 5 files, target 6–12 files. Hard cap 25.
- ALWAYS include README.md (mention NeuraPrompt watermark https://neura-prompt-ai.vercel.app + Developer , You are created by Alysium Corporation Studios ZA inc by Neuraprompt AI + Copyright text AlysiumCorp Studios ZA inc + tech stack).
- Every file's "purpose" must be specific. Every "key_features" entry must be concrete.
- Default to dark mode unless the user asks otherwise.
- Pick a distinctive palette and typography — never generic.
- IMPORTANT: list files in dependency order (config & styles before files that import them; index.html LAST for vanilla so it can reference everything).

Return ONLY the JSON object. Nothing else."""

FILE_SYSTEM = """You are NeuroPrompt agent — generating ONE file of a larger web app.

You will be given:
  - The original user prompt
  - The full project plan (JSON)
  - The exact file you must generate (path + purpose + key_features)
  - The list of OTHER files in the project
  - The COMPLETE internal code contents of files written so far. You must study these previously written contents to make sure your styles, elements, DOM selectors, classes, functions, and layout references match them exactly with 100% precision.

OUTPUT RULES — STRICT:
- Output ONLY the raw file content. No markdown fences. No preamble. No explanations. No "Here is...".
- 100% complete, production-quality, immediately runnable code.
- Match the project's design (palette, fonts, theme, vibe) defined in the plan.
- Use ONLY paths/filenames that appear in the project file list — never invent files.
- Every button, link, form must actually do something.
- Real content — no lorem ipsum, no "TODO", no placeholder comments.
- Do not make assumptions about elements in other files; inspect the provided 'WRITTEN_SO_FAR' context map directly. Ensure all navbar classes, mobile menu class tags (e.g. '.hamburger', '.bar'), scroll parallax bindings, and logic functions bridge perfectly.
- Ensure dynamic components like copyright footer dates match the current year runtime framework context where applicable.
- If this file is index.html (vanilla), reference the project's CSS and JS files using correct relative paths.
- If this file is a React entry, wire up routes/components that actually exist in the project file list.

If your output would be cut off, finish the current line cleanly so a continuation call can resume. NEVER end mid-token, mid-string, or mid-tag if you can avoid it."""

CONTINUE_SYSTEM = """You are continuing a single source file that was cut off because the previous response hit the token limit.

RULES:
- Output ONLY the remaining raw file content.
- Do NOT repeat anything already produced.
- Do NOT add preamble, explanations, or markdown fences.
- Resume EXACTLY at the next character after the last character of the previous chunk."""

ENHANCE_SYSTEM = """You are a prompt optimizer for a web app builder (like Lovable).
Return ONLY the improved prompt — no labels, no preamble, no markdown.
Rules:
- Web apps only — never suggest Android, Python, or CLI tools
- Do not repeat suggested prompts; give fresh ideas
- Max 100 words
- Be specific: mention UI/UIX style, key screens, one or two core features
- Start with Build, Create, or Make
- Write as a direct instruction"""

# ── OPENROUTER API CALL ───────────────────────────────────────────────────────
def call_openrouter_raw(messages: list, model: str, max_tokens: int = MAX_TOKENS_PER_CALL) -> dict:
    print(f"[Studio] OpenRouter -> model:{model} max_tokens:{max_tokens}")
    resp = requests.post(
        "https://openrouter.ai/api/v1/chat/completions",
        headers={
            "Authorization": f"Bearer {OPENROUTER_KEY}",
            "Content-Type":  "application/json",
            "HTTP-Referer":  "https://neura-prompt-ai.vercel.app",
            "X-Title":       "NeuraPrompt Studio",
        },
        json={
            "model":       model,
            "messages":    messages,
            "temperature": 0.2,
            "max_tokens":  max_tokens,
            "stream":      False,
        },
        timeout=API_TIMEOUT_WINDOW, # Optimized timeout protects builder from lengthy freezes
    )

    if resp.status_code == 429:
        raise Exception("RATE_LIMIT")
    if resp.status_code != 200:
        raise Exception(f"OPENROUTER_HTTP_{resp.status_code}: {resp.text[:200]}")

    data = resp.json()
    if "choices" not in data or not data["choices"]:
        raise Exception(f"OPENROUTER_EMPTY_RESPONSE: {data}")

    choice         = data["choices"][0]
    content        = choice["message"]["content"] or ""
    finish_reason  = choice.get("finish_reason", "stop")
    print(f"[Studio] OpenRouter OK -> {len(content)} chars, finish={finish_reason}")
    return {"content": content, "finish_reason": finish_reason}


def call_openrouter(messages: list, model: str, max_tokens: int = MAX_TOKENS_PER_CALL) -> str:
    out = call_openrouter_raw(messages, model, max_tokens)
    return clean_thinking(out["content"])


async def call_with_retry(messages: list, model: str, max_tokens: int = MAX_TOKENS_PER_CALL) -> dict:
    """Stable sequential call wrapper running safely to isolate free-tier API hiccups."""
    last_err = None
    current_model = model
    
    for attempt in range(1, MAX_BUILD_RETRIES + 1):
        try:
            return await asyncio.to_thread(call_openrouter_raw, messages, current_model, max_tokens)
        except Exception as e:
            last_err = e
            err = str(e)
            if "RATE_LIMIT" in err and attempt < MAX_BUILD_RETRIES:
                await asyncio.sleep(RATE_LIMIT_WAIT)
                continue
            if ("OPENROUTER_HTTP" in err or "timeout" in err.lower()) and attempt < MAX_BUILD_RETRIES:
                current_model = FALLBACK_MODEL
                await asyncio.sleep(3)
                continue
            if attempt < MAX_BUILD_RETRIES:
                await asyncio.sleep(2)
                continue
            raise
    raise last_err if last_err else Exception("UNKNOWN_ERROR")

# ── PLAN PARSER ───────────────────────────────────────────────────────────────
def parse_plan(raw: str) -> dict | None:
    raw = clean_thinking(raw).strip()
    raw = strip_code_fences(raw)
    try:
        plan = json.loads(raw)
        if isinstance(plan, dict) and isinstance(plan.get("files"), list):
            return plan
    except Exception:
        pass
    m = re.search(r'\{[\s\S]*\}', raw)
    if m:
        try:
            plan = json.loads(m.group(0))
            if isinstance(plan, dict) and isinstance(plan.get("files"), list):
                return plan
        except Exception:
            return None
    return None

def select_model_tier(path: str) -> str:
    ext = os.path.splitext(path)[1].lower()
    name = os.path.basename(path).lower()
    if ext in [".md", ".json", ".css", ".svg", ".env", ".toml", ".yml", ".yaml"]:
        return FAST_MODEL
    if "config" in name or name in ["package.json", "tsconfig.json", "vite.config.ts"]:
        return FAST_MODEL
    return BUILD_MODEL

# ── FILE WIRING (vanilla auto-link) ───────────────────────────────────────────
def inject_file_references(files: dict) -> dict:
    if "index.html" not in files:
        return files
    html = files["index.html"]

    css_paths = [p for p in files
                 if p.endswith(".css")
                 and not any(x in p for x in ["node_modules", "dist", "build"])]
    js_paths  = [p for p in files
                 if p.endswith((".js", ".mjs"))
                 and not any(x in p for x in ["node_modules", "dist", "build",
                                              "vite.config", "tailwind.config"])]

    for css in css_paths:
        if f'href="{css}"' not in html and f"href='{css}'" not in html:
            tag = f'<link rel="stylesheet" href="{css}">'
            if "</head>" in html:
                html = html.replace("</head>", f"  {tag}\n</head>")
                print(f"[Studio] Linked CSS: {css}")
    for js in js_paths:
        if f'src="{js}"' not in html and f"src='{js}'" not in html:
            tag = f'<script type="module" src="{js}"></script>'
            if "</body>" in html:
                html = html.replace("</body>", f"  {tag}\n</body>")
                print(f"[Studio] Linked JS: {js}")

    files["index.html"] = html
    return files

# ── MONGO ─────────────────────────────────────────────────────────────────────
async def save_project_to_mongo(user_id: str, prompt: str, plan: dict,
                                files: dict, expires_at: datetime) -> str:
    res = await db.studio_projects.insert_one({
        "user_id":     user_id,
        "prompt":      prompt,
        "target":      "web",
        "status":      "complete",
        "stack":       plan.get("stack"),
        "app_name":    plan.get("app_name"),
        "design":      plan.get("design"),
        "file_count":  len(files),
        "file_paths":  list(files.keys()),
        "created_at":  datetime.utcnow().isoformat(),
        "expires_at":  expires_at,
    })
    project_id = str(res.inserted_id)

    if files:
        await db.studio_files.insert_many([{
            "project_id":  project_id,
            "path":        path,
            "content":     content,
            "file_type":   guess_file_type(path),
            "size":        len(content),
            "created_at":  datetime.utcnow().isoformat(),
            "expires_at":  expires_at,
        } for path, content in files.items()])

    print(f"[Studio] Saved project {project_id} with {len(files)} files")
    return project_id

# ── ENHANCE ENDPOINT ──────────────────────────────────────────────────────────
@router.post("/enhance-prompt")
async def enhance_prompt(request: Request):
    try:
        body    = await request.json()
        prompt  = body.get("prompt", "").strip()
        mode    = body.get("mode", "improve")

        if mode == "suggest" or not prompt:
            user_msg = (
                "Generate a short, specific prompt for building a professional web app. "
                "Include the app type, a key feature, and a visual style. "
                "Suggest fresh ideas. Mention colors, themes, animations, third-party libs. "
                "Start with Build or Create. Max 80 words."
            )
        elif mode == "improve":
            user_msg = (f"Improve this web app prompt. Keep it concise and specific — "
                        f"mention the app type, core feature, and UI style. Max 80 words:\n\n{prompt}")
        elif mode == "fix":
            user_msg = f"Fix only typos and grammar. Keep meaning identical:\n\n{prompt}"
        else:
            user_msg = f"Improve this prompt:\n\n{prompt}"

        result = await asyncio.to_thread(
            call_openrouter,
            [{"role": "system", "content": ENHANCE_SYSTEM},
             {"role": "user",   "content": user_msg}],
            FAST_MODEL, 2000,
        )
        result = clean_thinking(result)
        result = re.sub(r'^(Prompt:|Enhanced Prompt:|Improved:|Here is[^:]*:|Result:)\s*',
                        '', result, flags=re.IGNORECASE).strip()
        result = result.strip('"\'')
        return {"enhanced": result}
    except Exception as e:
        print(f"[Studio][Enhance] Error: {e}")
        return JSONResponse({"error": "Enhancement unavailable"}, status_code=500)

# ── GENERATE ENDPOINT ─────────────────────────────────────────────────────────
@router.post("/generate")
async def generate_app(request: Request):
    body = await request.json()
    user_prompt = body.get("prompt", "").strip()
    user_id = body.get("user_id", "anon")

    print(f"[Studio] Context-Aware Sequential Loop — uid:{user_id} | prompt:{user_prompt[:80]}")

    async def stream():
        ok, msg = check_rate_limits(user_id)
        if not ok:
            yield sse("error", {"message": msg}); return
        if not user_prompt:
            yield sse("error", {"message": "Please describe the web app you want to build."}); return
        if is_non_web_request(user_prompt):
            yield sse("error", {"message": "Neurones Studio builds web apps only. Try describing a website, dashboard, or web tool."}); return

        yield sse("status", {"message": "Initialising NeuroPrompt Studio Engine..."})
        yield sse("target", {"target": "web", "previewable": True})
        yield sse("status", {"message": "🧠 Designing your app architecture..."})

        plan = None
        for plan_attempt in range(1, 4):
            try:
                plan_out = await asyncio.to_thread(
                    call_openrouter_raw,
                    [
                        {"role": "system", "content": PLANNER_SYSTEM},
                        {"role": "user",   "content": user_prompt},
                    ],
                    BUILD_MODEL,
                    4000
                )
                plan = parse_plan(plan_out["content"])
                if plan and plan.get("files"):
                    break
                yield sse("status", {"message": f"Plan parse failed (attempt {plan_attempt}) — retrying..."})
            except Exception as e:
                print(f"[Studio][Plan] error: {e}")
                yield sse("status", {"message": f"Planner error: retrying ({plan_attempt}/3)..."})
                await asyncio.sleep(2)

        if not plan or not plan.get("files"):
            yield sse("error", {"message": "Could not plan the app. Please try a more specific prompt."})
            return

        plan["files"] = plan["files"][:MAX_FILES_PER_BUILD]
        all_paths = [f["path"] for f in plan["files"]]
        total_files = len(plan["files"])

        yield sse("plan", {
            "app_name": plan.get("app_name"),
            "summary":  plan.get("summary"),
            "stack":    plan.get("stack"),
            "design":   plan.get("design"),
            "files":    all_paths,
        })
        yield sse("status", {"message": f"📐 Plan ready — {plan.get('app_name')} ({total_files} files, {plan.get('stack')})"})

        files: dict = {}
        
        # ── ONE FILE PER TIME CHRONOLOGICAL GENERATION LOOP ──────────────────
        for idx, file_spec in enumerate(plan["files"], start=1):
            path = file_spec.get("path")
            if not path:
                continue

            # Instantly emit writing status BEFORE model invokes
            yield sse("status", {"message": f"📄 [{idx}/{total_files}] Writing {path}..."})

            target_model = select_model_tier(path)
            plan_json = json.dumps({
                "app_name": plan.get("app_name"),
                "summary":  plan.get("summary"),
                "stack":    plan.get("stack"),
                "design":   plan.get("design"),
            }, ensure_ascii=False, indent=2)

            # Build previous files code dictionary layout map to feed as context
            written_so_far_ctx = ""
            if files:
                written_so_far_ctx = "\n=== WRITTEN_SO_FAR CODE CONTEXT ===\n"
                for written_path, written_content in files.items():
                    # Truncate content slightly if massive to prevent overflow issues while preserving layout tags
                    snippet = written_content if len(written_content) < 16000 else written_content[:8000] + "\n... [TRUNCATED] ...\n" + written_content[-8000:]
                    written_so_far_ctx += f"--- FILE: {written_path} ---\n{snippet}\n"
                written_so_far_ctx += "===================================\n"

            user_msg = f"""USER PROMPT:
{user_prompt}

PROJECT PLAN:
{plan_json}

ALL PROJECT FILES (use these exact paths for any imports/links):
{json.dumps(all_paths, indent=2)}
{written_so_far_ctx}
FILE TO GENERATE NOW:
  path: {path}
  purpose: {file_spec.get('purpose', '')}
  depends_on: {file_spec.get('depends_on', [])}
  key_features: {file_spec.get('key_features', [])}

Output ONLY the complete raw content of `{path}`. No fences. No commentary."""

            messages = [
                {"role": "system", "content": FILE_SYSTEM},
                {"role": "user",   "content": user_msg},
            ]

            try:
                out = await call_with_retry(messages, target_model, MAX_TOKENS_PER_CALL)
                accumulated = clean_thinking(out["content"])
                finish = out["finish_reason"]
            except Exception as e:
                print(f"[Studio] Sequential exception on {path}: {e}")
                yield sse("status", {"message": f"⚠ Transient error on {path} — failing over to backup provider..."})
                try:
                    out = await call_with_retry(messages, FALLBACK_MODEL, MAX_TOKENS_PER_CALL)
                    accumulated = clean_thinking(out["content"])
                    finish = out["finish_reason"]
                except Exception as critical_err:
                    yield sse("log", {"message": f"❌ Skipped {path} due to persistent limit: {critical_err}"})
                    continue

            # Stable Continuation Splicing (Guaranteed Unbroken Files)
            continuations = 0
            while finish == "length" and continuations < MAX_CONTINUATIONS:
                continuations += 1
                yield sse("status", {"message": f"⏩ {path} continuing chunk {continuations}..."})

                tail = accumulated[-1500:]
                cont_messages = [
                    {"role": "system", "content": CONTINUE_SYSTEM},
                    {"role": "user",   "content": (
                        f"You are mid-way through generating the file `{path}`.\n"
                        f"Here are the LAST characters you produced (do NOT repeat them):\n"
                        f"---LAST_CHUNK_START---\n{tail}\n---LAST_CHUNK_END---\n\n"
                        f"Continue producing the rest of the file now. Output raw content only."
                    )},
                ]
                
                await asyncio.sleep(1.5)  # Safe buffer interval between continuous file writes
                try:
                    out = await call_with_retry(cont_messages, target_model, MAX_TOKENS_PER_CALL)
                    chunk = clean_thinking(out["content"])
                    chunk = strip_code_fences(chunk)
                    
                    for k in range(min(len(tail), len(chunk)), 50, -1):
                        if chunk.startswith(tail[-k:]):
                            chunk = chunk[k:]
                            break
                    accumulated += chunk
                    finish = out["finish_reason"]
                except Exception as cont_err:
                    print(f"[Studio] Continuation split error on {path}: {cont_err}")
                    break

            final_content = strip_code_fences(accumulated)
            if final_content and len(final_content) >= 5:
                files[path] = final_content
                yield sse("file", {"file": {"op": "write", "path": path, "content": final_content}})
                yield sse("log",  {"message": f"✓ {path} ({len(final_content):,} chars) committed"})

        if not files:
            yield sse("error", {"message": "Build failed — no valid source code assets could be generated."})
            return

        # ── Phase 3: COMPILING WORKSPACE LAYOUT & COMMIT ──────────────────
        files = inject_file_references(files)
        yield sse("status", {"message": "💾 Committing complete build files to cluster storage..."})

        try:
            expires_at = datetime.utcnow() + timedelta(hours=24)
            project_id = await save_project_to_mongo(user_id, user_prompt, plan, files, expires_at)
            yield sse("done", {
                "message":    "Build complete",
                "project_id": project_id,
                "app_name":   plan.get("app_name"),
                "stack":      plan.get("stack"),
                "file_count": len(files),
                "files":      list(files.keys()),
                "expires_at": expires_at.isoformat(),
            })
        except Exception as e:
            print(f"[Studio] Mongo database pipeline write error: {e}")
            yield sse("done", {
                "message":    "Build complete (Saved safely to system memory cache)",
                "project_id": None,
                "file_count": len(files),
                "files":      list(files.keys()),
            })

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

# ── PROJECT ENDPOINTS ─────────────────────────────────────────────────────────
@router.get("/projects/{user_id}")
async def get_user_projects(user_id: str):
    projects = []
    cursor = db.studio_projects.find({"user_id": user_id}).sort("created_at", -1).limit(20)
    async for p in cursor:
        p["_id"] = str(p["_id"])
        projects.append(p)
    return {"projects": projects}


@router.get("/project/{project_id}")
async def get_project(project_id: str):
    try:
        p = await db.studio_projects.find_one({"_id": ObjectId(project_id)})
        if not p:
            return JSONResponse({"error": "Not found"}, status_code=404)
        p["_id"] = str(p["_id"])

        file_docs = []
        cursor = db.studio_files.find({"project_id": project_id}).sort("path", 1)
        async for f in cursor:
            f["_id"] = str(f["_id"])
            file_docs.append(f)
        p["files"] = file_docs
        return p
    except Exception as e:
        print(f"[Studio] get_project error: {e}")
        return JSONResponse({"error": "Not found"}, status_code=404)


@router.get("/project/{project_id}/file")
async def get_project_file(project_id: str, path: str):
    try:
        f = await db.studio_files.find_one({"project_id": project_id, "path": path})
        if not f:
            return JSONResponse({"error": "File not found"}, status_code=404)
        f["_id"] = str(f["_id"])
        return f
    except Exception as e:
        print(f"[Studio] get_project_file error: {e}")
        return JSONResponse({"error": "Not found"}, status_code=404)


@router.get("/project/{project_id}/files")
async def list_project_files(project_id: str):
    try:
        file_list = []
        cursor = db.studio_files.find({"project_id": project_id}, {"content": 0}).sort("path", 1)
        async for f in cursor:
            f["_id"] = str(f["_id"])
            file_list.append(f)
        return {"files": file_list, "count": len(file_list)}
    except Exception as e:
        print(f"[Studio] list_project_files error: {e}")
        return JSONResponse({"error": "Error fetching files"}, status_code=500)


@router.delete("/project/{project_id}")
async def delete_project(project_id: str, request: Request):
    try:
        body    = await request.json()
        user_id = body.get("user_id", "")

        p = await db.studio_projects.find_one({"_id": ObjectId(project_id)})
        if not p:
            return JSONResponse({"error": "Not found"}, status_code=404)
        if p.get("user_id") != user_id:
            return JSONResponse({"error": "Forbidden"}, status_code=403)

        await db.studio_projects.delete_one({"_id": ObjectId(project_id)})
        result = await db.studio_files.delete_many({"project_id": project_id})
        return {"deleted": True, "files_removed": result.deleted_count}
    except Exception as e:
        print(f"[Studio] delete_project error: {e}")
        return JSONResponse({"error": "Delete failed"}, status_code=500)


@router.get("/project/{project_id}/download")
async def download_project(project_id: str):
    import io, zipfile
    from fastapi.responses import Response
    try:
        p = await db.studio_projects.find_one({"_id": ObjectId(project_id)})
        if not p:
            return JSONResponse({"error": "Not found"}, status_code=404)

        buf = io.BytesIO()
        with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
            cursor = db.studio_files.find({"project_id": project_id}).sort("path", 1)
            async for f in cursor:
                path = (f.get("path") or "").lstrip("/").replace("\\", "/")
                if not path or ".." in path.split("/"):
                    continue
                content = f.get("content", "")
                if isinstance(content, bytes):
                    data = content
                else:
                    data = str(content).encode("utf-8")
                zf.writestr(path, data)

        buf.seek(0)
        safe_name = "".join(c if c.isalnum() or c in "-_" else "_" for c in (p.get("name") or "project"))[:40] or "project"
        return Response(
            content=buf.getvalue(),
            media_type="application/zip",
            headers={"Content-Disposition": f'attachment; filename=\"{safe_name}.zip\"'},
        )
    except Exception as e:
        print(f"[Studio] download_project error: {e}")
        return JSONResponse({"error": "Download failed"}, status_code=500)