"""
NeuraPrompt Studio v2.0 — Credit-based autonomous web-app builder
==================================================================

Changes over v1:
  * Credit system replaces token usage (5 free/month, 10 pro/month)
  * Project statistics tracking (MongoDB)
  * Polar multi-plan via polar_subs_studio (free/pro/max/business credits + features)
  * Firebase auth email resolution
  * Speed optimizations (reduced max_tokens, max_files, continuations)
  * /plan endpoint fixed (uses call_openrouter_raw for proper response)
  * /fix endpoint for bug fixes on existing projects
  * /credits endpoint for credit status
  * /stats endpoint for global statistics

Endpoints:
  POST /api/v1/studio/generate     — Build a complete web app (costs 1 credit)
  POST /api/v1/studio/plan         — Plan-chat with AI (costs 1 credit per session)
  POST /api/v1/studio/fix          — Fix/update existing project (costs 1 credit)
  POST /api/v1/studio/enhance-prompt — Enhance a prompt (free, no credit)
  GET  /api/v1/studio/credits/{uid} — Get credit status
  GET  /api/v1/studio/usage/{uid}   — Get token usage stats (for display)
  GET  /api/v1/studio/stats         — Global statistics (admin)
  GET  /api/v1/studio/projects/{uid} — List user's projects
  GET  /api/v1/studio/project/{id}   — Get project + files
  GET  /api/v1/studio/project/{id}/download — Download as ZIP
  DELETE /api/v1/studio/project/{id}  — Delete project
"""

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

from fastapi import APIRouter, Request, Form, HTTPException, Query, UploadFile, File
from fastapi.responses import StreamingResponse, JSONResponse, Response
from pymongo import MongoClient
from pymongo.server_api import ServerApi
from bson import ObjectId
from pydantic import BaseModel

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

# ── ENV ───────────────────────────────────────────────────────────────────────
OPENROUTER_KEY = os.environ.get("OPENROUTE_KEY", "")
MONGO_URL      = os.environ.get("MONGO_URI", "")
GROQ_API_KEY   = os.environ.get("GROQ_API_KEY", "")
POLAR_API_KEY  = os.environ.get("POLAR_API_KEY", "")
POLAR_ORG_ID   = os.environ.get("POLAR_ORG_ID", "")
FIREBASE_KEY_PATH = os.environ.get("FIREBASE_KEY_PATH", "serviceAccountKey.json")
ADMIN_UIDS     = set(os.environ.get("STUDIO_ADMIN_UIDS", "").split(","))

# ── MONGODB (lazy — waits until main.py is loaded, then reuses its connection) ─
import sys as _sys

_db_instance = None

def _get_db():
    """Get the neuraprompt database. LAZY: reuses main.py's connection.
    Called on first DB operation, not at import time — this ensures main.py
    has finished initializing its MongoClient before we try to reuse it."""
    global _db_instance
    if _db_instance is not None:
        return _db_instance
    # Try to reuse main.py's connection (it's already connected to Atlas)
    main_mod = _sys.modules.get("main")
    if main_mod and hasattr(main_mod, "neuraprompt_db"):
        _db_instance = main_mod.neuraprompt_db
        print("[Studio v2] MongoDB: reusing main.py connection")
        return _db_instance
    # Fallback: create our own pymongo client (same SSL settings as main.py)
    print("[Studio v2] MongoDB: main.py not available, creating own connection")
    _client = MongoClient(
        MONGO_URL, ssl=True,
        tlsAllowInvalidCertificates=False,
        tlsCAFile="/etc/ssl/certs/ca-certificates.crt",
        server_api=ServerApi("1"),
    )
    _db_instance = _client["neuraprompt"]
    return _db_instance

# NOTE: db is NOT set at import time — it's resolved lazily on first use.
# All _db_* helpers call _get_db() internally.

# ── ASYNC DB HELPER (pymongo is sync, so we wrap calls in asyncio.to_thread) ──
async def _db_find_one(collection, query, projection=None):
    """Async wrapper for db.collection.find_one()"""
    return await asyncio.to_thread(
        lambda: _get_db()[collection].find_one(query, projection) if projection else _get_db()[collection].find_one(query)
    )

async def _db_update_one(collection, query, update, upsert=False):
    return await asyncio.to_thread(lambda: _get_db()[collection].update_one(query, update, upsert=upsert))

async def _db_insert_one(collection, doc):
    return await asyncio.to_thread(lambda: _get_db()[collection].insert_one(doc))

async def _db_count_documents(collection, query):
    return await asyncio.to_thread(lambda: _get_db()[collection].count_documents(query))

async def _db_find(collection, query, projection=None, sort=None, limit=None):
    def _do():
        cursor = _get_db()[collection].find(query, projection) if projection else _get_db()[collection].find(query)
        if sort: cursor = cursor.sort(*sort) if isinstance(sort, tuple) else cursor.sort(sort)
        if limit: cursor = cursor.limit(limit)
        return list(cursor)
    return await asyncio.to_thread(_do)

async def _db_delete_one(collection, query):
    return await asyncio.to_thread(lambda: _get_db()[collection].delete_one(query))

async def _db_delete_many(collection, query):
    return await asyncio.to_thread(lambda: _get_db()[collection].delete_many(query))

async def _db_insert_many(collection, docs):
    return await asyncio.to_thread(lambda: _get_db()[collection].insert_many(docs))

async def _db_distinct(collection, field, query=None):
    return await asyncio.to_thread(lambda: _get_db()[collection].distinct(field, query or {}))

async def _db_aggregate(collection, pipeline):
    return await asyncio.to_thread(lambda: list(_get_db()[collection].aggregate(pipeline)))

async def _db_create_index(collection, keys, **kwargs):
    return await asyncio.to_thread(lambda: _get_db()[collection].create_index(keys, **kwargs))


# ── CREDIT SYSTEM (aligned with polar_subs_studio entitlements) ───────────────
# Monthly build credits = plan.builds from Studio Polar entitlements.
# free: 5 | pro: 5 | max: 10 | business: 15
FREE_CREDITS = 5
PLAN_CREDITS = {
    "free": 5,
    "pro": 15,
    "max": 20,
    "business": 25,
}
# Monthly AI token budgets (prompt+completion) — displayed + soft-capped in usage
PLAN_TOKENS = {
    "free": 5_000,
    "pro": 100_000,
    "max": 1_000_000,
    "business": 1_000_000,
}
# Legacy alias used in error messages / usage payload
PRO_CREDITS = PLAN_CREDITS["max"]

# Feature flags by plan (mirror polar_subs_studio.STUDIO_ENTITLEMENTS)
PLAN_FEATURES = {
    "free": {"canFix": False, "canPlan": False, "canDownload": False, "canBuild": True},
    "pro": {"canFix": False, "canPlan": False, "canDownload": True, "canBuild": True},
    "max": {"canFix": True, "canPlan": True, "canDownload": True, "canBuild": True},
    "business": {"canFix": True, "canPlan": True, "canDownload": True, "canBuild": True},
}


def _month_key():
    return datetime.now(timezone.utc).strftime("%Y-%m")


def _credits_for_plan(plan: str) -> int:
    return PLAN_CREDITS.get((plan or "free").lower(), FREE_CREDITS)


def _features_for_plan(plan: str) -> dict:
    return dict(PLAN_FEATURES.get((plan or "free").lower(), PLAN_FEATURES["free"]))


async def get_credits(user_id: str, email: str = "") -> dict:
    """Get user's credit status. Tier comes from polar_subs_studio (pro/max/business/free)."""
    month = _month_key()
    plan = "free"
    entitlements: dict = {}
    try:
        plan_result = await asyncio.wait_for(
            _check_studio_plan(user_id, email=email),
            timeout=5.0,
        )
        plan = plan_result.get("plan") or "free"
        entitlements = plan_result.get("entitlements") or _features_for_plan(plan)
    except (asyncio.TimeoutError, Exception) as e:
        print(f"[Studio] Studio Polar check timed out or failed ({e}) — defaulting to free tier")
        plan = "free"
        entitlements = _features_for_plan("free")

    plan = plan if plan in PLAN_CREDITS else "free"
    target_total = _credits_for_plan(plan)
    features = entitlements if entitlements else _features_for_plan(plan)

    doc = await _db_find_one("studio_credits", {"user_id": user_id, "month_key": month})
    if doc:
        remaining = int(doc.get("credits_remaining", target_total))
        total = int(doc.get("credits_total", target_total))
        prev_plan = (doc.get("tier") or doc.get("plan") or "free").lower()

        # Upgrade: grant the difference when moving to a higher credit allowance
        if target_total > total:
            remaining += (target_total - total)
            total = target_total
            await _db_update_one(
                "studio_credits",
                {"user_id": user_id, "month_key": month},
                {"$set": {
                    "credits_remaining": remaining,
                    "credits_total": total,
                    "tier": plan,
                    "plan": plan,
                }},
            )
        # Downgrade or plan rename: clamp remaining to new total
        elif target_total < total or prev_plan != plan:
            total = target_total
            remaining = min(remaining, target_total)
            await _db_update_one(
                "studio_credits",
                {"user_id": user_id, "month_key": month},
                {"$set": {
                    "credits_remaining": remaining,
                    "credits_total": total,
                    "tier": plan,
                    "plan": plan,
                }},
            )

        tokens_used = int(doc.get("tokens_used") or 0)
        tokens_monthly = _tokens_budget(plan)
        return {
            "remaining": remaining,
            "total": total,
            "tier": plan,
            "plan": plan,
            "month": month,
            "features": features,
            "entitlements": features,
            "tokens_used": tokens_used,
            "tokens_monthly": tokens_monthly,
            "tokens_input": int(doc.get("tokens_input") or 0),
            "tokens_output": int(doc.get("tokens_output") or 0),
        }

    # New month — create record
    tokens_monthly = _tokens_budget(plan)
    await _db_insert_one("studio_credits", {
        "user_id": user_id,
        "month_key": month,
        "credits_remaining": target_total,
        "credits_total": target_total,
        "tier": plan,
        "plan": plan,
        "tokens_used": 0,
        "tokens_input": 0,
        "tokens_output": 0,
        "history": [],
        "created_at": datetime.now(timezone.utc),
    })
    return {
        "remaining": target_total,
        "total": target_total,
        "tier": plan,
        "plan": plan,
        "month": month,
        "features": features,
        "entitlements": features,
        "tokens_used": 0,
        "tokens_monthly": tokens_monthly,
        "tokens_input": 0,
        "tokens_output": 0,
    }


async def consume_credit(user_id: str, action: str, project_id: str = "") -> bool:
    """Consume 1 credit. Returns True if successful, False if no credits."""
    month = _month_key()
    result = await _db_update_one(
        "studio_credits",
        {"user_id": user_id, "month_key": month, "credits_remaining": {"$gt": 0}},
        {
            "$inc": {"credits_remaining": -1},
            "$push": {
                "history": {
                    "action": action,
                    "project_id": project_id,
                    "timestamp": datetime.now(timezone.utc).isoformat(),
                }
            },
        },
    )
    return result.modified_count > 0



async def record_tokens(user_id: str, input_tokens: int = 0, output_tokens: int = 0) -> None:
    """Accumulate OpenRouter tokens for the current month (server truth)."""
    if not user_id or user_id == "anon":
        return
    total = int(input_tokens or 0) + int(output_tokens or 0)
    if total <= 0:
        return
    month = _month_key()
    try:
        await _db_update_one(
            "studio_credits",
            {"user_id": user_id, "month_key": month},
            {
                "$inc": {
                    "tokens_used": total,
                    "tokens_input": int(input_tokens or 0),
                    "tokens_output": int(output_tokens or 0),
                },
                "$setOnInsert": {
                    "credits_remaining": FREE_CREDITS,
                    "credits_total": FREE_CREDITS,
                    "tier": "free",
                    "plan": "free",
                    "history": [],
                    "created_at": datetime.now(timezone.utc),
                },
            },
            upsert=True,
        )
    except Exception as e:
        print(f"[Studio] record_tokens error: {e}")


def _tokens_budget(plan: str) -> int:
    return PLAN_TOKENS.get((plan or "free").lower(), PLAN_TOKENS["free"])

async def check_credits(user_id: str, email: str = "") -> tuple:
    """Returns (has_credits, remaining, total, tier/plan)."""
    status = await get_credits(user_id, email=email)
    return status["remaining"] > 0, status["remaining"], status["total"], status["tier"]


# ── POLAR (Studio multi-plan via polar_subs_studio) ────────────────────────────
_studio_polar_mod = None


def _get_studio_polar():
    """Lazy-load polar_subs_studio (sibling module or package import)."""
    global _studio_polar_mod
    if _studio_polar_mod is not None:
        return _studio_polar_mod
    try:
        import polar_subs_studio as _ps
        _studio_polar_mod = _ps
        return _ps
    except ImportError:
        try:
            from . import polar_subs_studio as _ps
            _studio_polar_mod = _ps
            return _ps
        except ImportError as e:
            print(f"[Studio] polar_subs_studio not available: {e}")
            return None


async def _check_studio_plan(user_id: str, email: str = "") -> dict:
    """
    Resolve Studio plan via polar_subs_studio.
    Returns dict with plan, entitlements, subscribed, matched_by, ...
    Fail-closed → free.
    """
    ps = _get_studio_polar()
    if ps is None:
        is_paid = await _check_polar_legacy(user_id)
        plan = "pro" if is_paid else "free"
        return {
            "plan": plan,
            "subscribed": is_paid,
            "entitlements": _features_for_plan(plan),
            "matched_by": "legacy_fallback",
        }

    try:
        if not email:
            email = await _resolve_email(user_id)
        subs_col = None
        try:
            main_mod = _sys.modules.get("main")
            if main_mod and hasattr(main_mod, "subscriptions_col"):
                subs_col = main_mod.subscriptions_col
        except Exception:
            pass

        result = await ps.check_studio_plan(
            email=email or "",
            firebase_uid=user_id or "",
            subscriptions_col=subs_col,
            fail_open_on_outage=False,
        )
        data = result.to_dict() if hasattr(result, "to_dict") else {
            "plan": getattr(result, "plan", "free"),
            "subscribed": getattr(result, "subscribed", False),
            "entitlements": getattr(result, "entitlements", None) or _features_for_plan("free"),
            "matched_by": getattr(result, "matched_by", ""),
            "status": getattr(result, "status", ""),
        }
        plan = (data.get("plan") or "free").lower()
        if plan not in PLAN_CREDITS:
            plan = "free"
        data["plan"] = plan
        data["entitlements"] = data.get("entitlements") or _features_for_plan(plan)
        print(
            f"[Studio] Polar plan uid={user_id}: plan={plan} "
            f"matched_by={data.get('matched_by')} status={data.get('status')}"
        )
        return data
    except Exception as e:
        print(f"[Studio] polar_subs_studio error: {e}")
        return {
            "plan": "free",
            "subscribed": False,
            "entitlements": _features_for_plan("free"),
            "matched_by": "error",
            "error": str(e),
        }


async def _check_polar_legacy(user_id: str) -> bool:
    """Binary subscribed check (Kype polar_subscription or Mongo). Fail-closed."""
    try:
        main_mod = _sys.modules.get("main")
        if main_mod and hasattr(main_mod, "verify_polar_subscription"):
            email = await _resolve_email(user_id)
            result = await asyncio.wait_for(
                main_mod.verify_polar_subscription(email=email, firebase_uid=user_id),
                timeout=4.0,
            )
            return bool(result)
    except asyncio.TimeoutError:
        print(f"[Studio] Legacy Polar check timed out for {user_id}")
    except Exception as e:
        print(f"[Studio] Legacy Polar check error: {e}")
    try:
        doc = await _db_find_one("subscriptions", {"user_id": user_id})
        if doc and doc.get("status") == "active" and doc.get("tier") in (
            "pro", "max", "business", "ultra", "premium"
        ):
            return True
    except Exception:
        pass
    return False


async def _check_polar(user_id: str) -> bool:
    """Backward-compatible bool: True if plan is any paid Studio tier."""
    data = await _check_studio_plan(user_id)
    return bool(data.get("subscribed")) or (data.get("plan") or "free") != "free"


# ── FIREBASE EMAIL RESOLUTION ─────────────────────────────────────────────────
_firebase_inited = False
_fb_auth = None

def _init_firebase():
    global _firebase_inited, _fb_auth
    if _firebase_inited:
        return _fb_auth is not None
    _firebase_inited = True
    try:
        import firebase_admin
        from firebase_admin import auth as _auth
        try:
            firebase_admin.get_app()
        except ValueError:
            if os.path.exists(FIREBASE_KEY_PATH):
                from firebase_admin import credentials
                cred = credentials.Certificate(FIREBASE_KEY_PATH)
                firebase_admin.initialize_app(cred)
            else:
                return False
        _fb_auth = _auth
        return True
    except Exception as e:
        print(f"[Studio] Firebase init failed: {e}")
        return False

async def _resolve_email(user_id: str) -> str:
    """Resolve Firebase UID to email. 3s timeout — don't block on Firebase."""
    if not user_id or user_id == "anon":
        return ""
    if not _init_firebase():
        return ""
    try:
        loop = asyncio.get_event_loop()
        user = await asyncio.wait_for(
            loop.run_in_executor(None, _fb_auth.get_user, user_id),
            timeout=3.0
        )
        return user.email or ""
    except asyncio.TimeoutError:
        print(f"[Studio] Firebase email resolution timed out for {user_id}")
        return ""
    except Exception as e:
        print(f"[Studio] Firebase email resolution failed for {user_id}: {e}")
        return ""

# ── PROJECT STATISTICS ────────────────────────────────────────────────────────
async def _update_stats():
    """Update global statistics after a build."""
    try:
        total_projects = await _db_count_documents("studio_projects", {})
        total_users = await _db_distinct("studio_projects", "user_id")
        now = datetime.now(timezone.utc)
        today_start = now.replace(hour=0, minute=0, second=0, microsecond=0)
        month_start = now.replace(day=1, hour=0, minute=0, second=0, microsecond=0)
        week_ago = now - timedelta(days=7)

        projects_today = await _db_count_documents("studio_projects", {"created_at": {"$gte": today_start.isoformat()}})
        projects_this_month = await _db_count_documents("studio_projects", {"created_at": {"$gte": month_start.isoformat()}})
        active_users_7d = await _db_distinct("studio_projects", "user_id", {"created_at": {"$gte": week_ago.isoformat()}})

        # Stack distribution
        stack_pipeline = [
            {"$group": {"_id": "$stack", "count": {"$sum": 1}}},
        ]
        stack_data = await _db_aggregate("studio_projects", stack_pipeline)
        stack_dist = {s["_id"] or "unknown": s["count"] for s in stack_data}

        await _db_update_one("studio_stats", 
            {"_id": "global"},
            {"$set": {
                "total_projects": total_projects,
                "total_users": len(total_users),
                "active_users_7d": len(active_users_7d),
                "projects_today": projects_today,
                "projects_this_month": projects_this_month,
                "stack_distribution": stack_dist,
                "last_updated": datetime.now(timezone.utc),
            }},
            upsert=True,
        )
    except Exception as e:
        print(f"[Studio] Stats update error: {e}")

# ── MODELS ────────────────────────────────────────────────────────────────────
BUILD_MODEL     = "nvidia/nemotron-3.5-super-vl:free"  # Fast model (was tencent/hy3:free — too slow)
STRONG_MODEL    = "nvidia/nemotron-3.5-super-vl:free"              # Strong model (for index.html only)
FAST_MODEL      = "liquid/lfm-2.5-2.6b:free"
FALLBACK_MODEL  = "liquid/lfm-2.5-2.6b:free"
PLAN_MODEL      = "liquid/lfm-2.5-2.6b:free"
FIX_MODEL       = "nvidia/nemotron-3.5-super-vl:free"

# ── SPEED OPTIMIZATIONS ───────────────────────────────────────────────────────
MAX_TOKENS_PER_CALL  = 8000    # Was 18000 — most web files are under 4K tokens
RATE_LIMIT_WAIT      = 15      # Was 45 — enough for OpenRouter free tier
MAX_BUILD_RETRIES    = 2       # Was 3 — fail faster
MAX_CONTINUATIONS    = 3       # Was 6 — most files complete in 1 call
MAX_FILES_FREE       = 8       # Free tier: max 8 files per build
MAX_FILES_PRO        = 15      # Pro tier: max 15 files per build
API_TIMEOUT_WINDOW   = 45      # Keep — needed for large files

# ── TTL INDEX ─────────────────────────────────────────────────────────────────
@router.on_event("startup")
async def startup():
    try:
        await _db_create_index("studio_projects", "expires_at", expireAfterSeconds=0)
        await _db_create_index("studio_files", "project_id")
        await _db_create_index("studio_files", "expires_at", expireAfterSeconds=0)
        await _db_create_index("studio_credits", [("user_id", 1), ("month_key", 1)], unique=True)
        print("[Studio v2] Indexes ready — credits, projects, files")
    except Exception as e:
        print(f"[Studio v2] 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()
    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. Try again in {int(wait)+1}s."
    ok, wait = user_limiters[user_id].is_allowed()
    if not ok:
        return False, f"Building too fast. Wait {int(wait)+1}s."
    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"}.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_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)

def _extract_tokens(data: dict) -> tuple:
    usage = data.get("usage") or {}
    return int(usage.get("prompt_tokens", 0) or 0), int(usage.get("completion_tokens", 0) or 0)

# ── 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",
  "stack": "html" | "vanilla" | "react-vite",
  "design": {"theme":"dark"|"light","palette":["#hex"],"font_pair":"...","vibe":"..."},
  "files": [{"path":"index.html","purpose":"...","depends_on":[],"key_features":[]}]
}

PLANNING RULES:
- Build a REAL product, not a demo.
- For vanilla: index.html, styles.css, script.js, README.md.
- For html: index.html, styles.css, script.js, README.md.
- Minimum 4 files, target 5-8 files. Hard cap 12.
- ALWAYS include README.md (mention NeuraPrompt watermark https://neura-prompt-ai.vercel.app + Developer: Alysium Corporation Studios ZA inc + Copyright AlysiumCorp Studios ZA inc + tech stack).
- Default to dark mode unless user asks otherwise.
- List files in dependency order (styles before index.html).
Return ONLY the JSON object."""

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

OUTPUT RULES — STRICT:
- Output ONLY the raw file content. No markdown fences. No preamble. No explanations.
- 100% complete, production-quality, immediately runnable code.
- Match the project's design (palette, fonts, theme) defined in the plan.
- Real content — no lorem ipsum, no TODO, no placeholder comments.
- If this file is index.html, reference the project's CSS and JS files using correct relative paths.

If your output would be cut off, finish the current line cleanly so a continuation call can resume."""

CONTINUE_SYSTEM = """You are continuing a single source file that was cut off.
RULES:
- Output ONLY the remaining raw file content.
- Do NOT repeat anything already produced.
- 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.
Return ONLY the improved prompt — no labels, no preamble, no markdown.
Rules: Web apps only. Max 100 words. Be specific. Start with Build, Create, or Make."""

PLAN_CHAT_SYSTEM = """You are NeuraPrompt Studio's planning assistant. You help users refine their app ideas BEFORE building. You ask clarifying questions, suggest features, and help them think through the architecture. Be concise.

RULES:
- Keep responses short (2-4 paragraphs max)
- Ask ONE clarifying question at a time
- Suggest specific features, libraries, or UI patterns
- If the user's idea is vague, help them narrow it down
- Don't write code — that's for the build phase"""

FIX_SYSTEM = """You are NeuraPrompt Studio's code fixer. You receive a project's files and a bug report or update request. You output ONLY the updated files as JSON.

OUTPUT FORMAT — RETURN ONLY VALID JSON:
{
  "files": [
    {"path": "index.html", "content": "complete updated file content", "action": "update"},
    {"path": "new-file.js", "content": "...", "action": "create"}
  ],
  "summary": "What you changed and why"
}

RULES:
- Output ONLY the files that changed (or new files). Don't include unchanged files.
- Each file's content must be COMPLETE — not a diff, the full file.
- If you can't fix it, return {"files": [], "summary": "Could not fix — reason"}."""

# ── 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,
    )
    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")
    choice = data["choices"][0]
    content = choice["message"]["content"] or ""
    finish_reason = choice.get("finish_reason", "stop")
    inp, out = _extract_tokens(data)
    print(f"[Studio] OpenRouter OK -> {len(content)} chars, finish={finish_reason}, tokens: {inp}+{out}")
    return {"content": content, "finish_reason": finish_reason, "usage": {"input_tokens": inp, "output_tokens": out}}

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:
    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: 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: pass
    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
    if name == "index.html":
        return STRONG_MODEL  # Use strong model for main HTML only
    return BUILD_MODEL

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>")
    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>")
    files["index.html"] = html
    return files

# ── MONGO SAVE ────────────────────────────────────────────────────────────────
async def save_project_to_mongo(user_id: str, prompt: str, plan: dict, files: dict, expires_at: datetime, tier: str = "free") -> str:
    email = await _resolve_email(user_id)
    res = await _db_insert_one("studio_projects", {
        "user_id": user_id, "user_email": email, "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()),
        "tier": tier,
        "created_at": datetime.utcnow().isoformat(),
        "expires_at": expires_at,
    })
    project_id = str(res.inserted_id)
    if files:
        await _db_insert_many("studio_files", [{
            "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} ({len(files)} files, tier={tier})")
    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. Max 80 words. Start with Build or Create."
        elif mode == "improve":
            user_msg = f"Improve this web app prompt. Keep it concise and specific. Max 80 words:\n\n{prompt}"
        elif mode == "fix":
            user_msg = f"Fix only typos and grammar:\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:
        err = str(e)
        if "RATE_LIMIT" in err:
            return JSONResponse({"error": "Rate limited. Please wait a moment and try again."}, status_code=429)
        print(f"[Studio][Enhance] Error: {e}")
        return JSONResponse({"error": "Enhancement unavailable"}, status_code=500)

# ── PLAN-CHAT ENDPOINT (uses main.py's /chat/ for Groq fallback) ──────────────
PLAN_CHAT_SYSTEM = """You are NeuraPrompt Studio's planning assistant. You help users refine their app ideas BEFORE building. You ask clarifying questions, suggest features, and help them think through the architecture. Be concise.

RULES:
- Keep responses short (2-4 paragraphs max)
- Ask ONE clarifying question at a time
- Suggest specific features, libraries, or UI patterns
- If the user's idea is vague, help them narrow it down
- Don't write code — that's for the build phase"""

@router.post("/plan")
async def plan_chat(request: Request):
    """Plan-chat: discuss the app idea with AI before building.
    Uses main.py's /chat/ endpoint for automatic Groq fallback."""
    try:
        body = await request.json()
        user_id = body.get("user_id", "anon")
        message = body.get("message", "").strip()
        history = body.get("history", [])
        if not message:
            return JSONResponse({"error": "Message required"}, status_code=400)

        # Plan feature + credits (Max/Business only for /plan chat)
        credits_info = await get_credits(user_id)
        features = credits_info.get("features") or _features_for_plan(credits_info["tier"])
        if not features.get("canPlan"):
            return JSONResponse({
                "error": "Planning before building is a Max feature. Upgrade to Max or Business.",
                "credits": {
                    "remaining": credits_info["remaining"],
                    "total": credits_info["total"],
                    "tier": credits_info["tier"],
                    "plan": credits_info.get("plan", credits_info["tier"]),
                    "features": features,
                },
            }, status_code=402)
        has_credits, remaining, total, tier = (
            credits_info["remaining"] > 0,
            credits_info["remaining"],
            credits_info["total"],
            credits_info["tier"],
        )
        if not has_credits:
            return JSONResponse({"error": "No credits remaining this month", "credits": {"remaining": 0, "total": total, "tier": tier, "plan": tier, "features": features}}, status_code=402)

        # Build conversation context from history
        conversation_context = ""
        for msg in history[-10:]:
            role = "User" if msg.get("role") == "user" else "AI"
            conversation_context += f"{role}: {msg.get('content', '')}\n"

        # Use main.py's /chat/ endpoint — it has Groq fallback built in
        # This avoids OpenRouter rate limit issues that break direct calls
        chat_body = {
            "user_id": user_id,
            "message": f"[STUDIO PLAN MODE]\nConversation so far:\n{conversation_context}\n\nUser's new message: {message}\n\nRespond as the planning assistant. Help refine the app idea.",
            "model_id": "neurones-flash-2.0",
            "persona": "default",
            "deep_think": False,
            "deep_search": False,
            "instructions": PLAN_CHAT_SYSTEM,
        }

        # Call main.py's /chat/ endpoint (same server, internal call)
        import httpx
        async with httpx.AsyncClient(timeout=30.0) as client:
            r = await client.post(
                "http://localhost:7860/chat/",
                json=chat_body,
                headers={"Content-Type": "application/json"},
            )

        if r.status_code != 200:
            # Fallback: try OpenRouter directly if /chat/ fails
            print(f"[Studio][Plan] /chat/ returned {r.status_code}, falling back to OpenRouter")
            messages = [{"role": "system", "content": PLAN_CHAT_SYSTEM}]
            for msg in history[-10:]:
                messages.append({"role": msg.get("role", "user"), "content": msg.get("content", "")})
            messages.append({"role": "user", "content": message})
            try:
                out = await asyncio.wait_for(
                    asyncio.to_thread(call_openrouter_raw, messages, PLAN_MODEL, 2000),
                    timeout=30.0
                )
                result = clean_thinking(out["content"])
            except:
                return JSONResponse({"error": "AI unavailable. Try again."}, status_code=503)
        else:
            data = r.json()
            result = data.get("response", "")
            if not result:
                return JSONResponse({"error": "Empty response from AI"}, status_code=500)

        return {"response": result, "credits": {"remaining": remaining, "total": total, "tier": tier}}
    except asyncio.TimeoutError:
        print(f"[Studio][Plan] Timed out")
        return JSONResponse({"error": "AI is taking too long. Try again."}, status_code=504)
    except Exception as e:
        print(f"[Studio][Plan] Error: {e}")
        return JSONResponse({"error": f"Plan failed: {e}"}, status_code=500)


# ── FIX/UPDATE ENDPOINT (uses main.py's /chat/ for Groq fallback) ────────────
FIX_SYSTEM = """You are NeuraPrompt Studio's code fixer. You receive a project's files and a bug report or update request. You output ONLY the updated files as JSON.

OUTPUT FORMAT — RETURN ONLY VALID JSON:
{
  "files": [
    {"path": "index.html", "content": "complete updated file content", "action": "update"},
    {"path": "new-file.js", "content": "...", "action": "create"}
  ],
  "summary": "What you changed and why"
}

RULES:
- Output ONLY the files that changed (or new files). Don't include unchanged files.
- Each file's content must be COMPLETE — not a diff, the full file.
- If you can't fix it, return {"files": [], "summary": "Could not fix — reason"}"""

@router.post("/fix")
async def fix_project(request: Request):
    """Fix bugs or update an existing project. Costs 1 credit.
    Uses main.py's /chat/ endpoint for automatic Groq fallback."""
    try:
        body = await request.json()
        user_id = body.get("user_id", "anon")
        project_id = body.get("project_id", "")
        instruction = body.get("instruction", "").strip()
        files = body.get("files", {})
        if not instruction:
            return JSONResponse({"error": "Instruction required"}, status_code=400)
        if not files:
            return JSONResponse({"error": "No files to fix"}, status_code=400)

        # Fix is Max/Business only + credits
        credits_info = await get_credits(user_id)
        features = credits_info.get("features") or _features_for_plan(credits_info["tier"])
        if not features.get("canFix"):
            return JSONResponse({
                "error": "Fix bugs is a Max feature. Upgrade to Max or Business.",
                "credits": {
                    "remaining": credits_info["remaining"],
                    "total": credits_info["total"],
                    "tier": credits_info["tier"],
                    "plan": credits_info.get("plan", credits_info["tier"]),
                    "features": features,
                },
            }, status_code=402)
        has_credits, remaining, total, tier = (
            credits_info["remaining"] > 0,
            credits_info["remaining"],
            credits_info["total"],
            credits_info["tier"],
        )
        if not has_credits:
            return JSONResponse({"error": "No credits remaining", "credits": {"remaining": 0, "total": total, "tier": tier, "plan": tier, "features": features}}, status_code=402)
        consumed = await consume_credit(user_id, "fix", project_id)
        if not consumed:
            return JSONResponse({"error": "Could not consume credit"}, status_code=402)

        file_context = ""
        for path, content in files.items():
            snippet = content if len(content) < 8000 else content[:4000] + "\n... [TRUNCATED] ...\n" + content[-4000:]
            file_context += f"--- FILE: {path} ---\n{snippet}\n\n"

        fix_prompt = f"PROJECT FILES:\n{file_context}\nFIX/UPDATE REQUEST:\n{instruction}\n\nOutput the updated/new files as JSON. Only include files that changed."

        # Use main.py's /chat/ endpoint — it has Groq fallback built in
        chat_body = {
            "user_id": user_id,
            "message": fix_prompt,
            "model_id": "neurones-pro-1.0",
            "persona": "engineer",
            "deep_think": True,
            "deep_search": False,
            "instructions": FIX_SYSTEM,
        }

        import httpx
        result = ""
        try:
            async with httpx.AsyncClient(timeout=60.0) as client:
                r = await client.post(
                    "http://localhost:7860/chat/",
                    json=chat_body,
                    headers={"Content-Type": "application/json"},
                )
            if r.status_code == 200:
                data = r.json()
                result = data.get("response", "")
            else:
                print(f"[Studio][Fix] /chat/ returned {r.status_code}, falling back to OpenRouter")
                raise Exception("chat_failed")
        except Exception:
            # Fallback: try OpenRouter directly
            try:
                result = await asyncio.wait_for(
                    asyncio.to_thread(call_openrouter,
                        [{"role": "system", "content": FIX_SYSTEM},
                         {"role": "user", "content": fix_prompt}],
                        FIX_MODEL, 12000),
                    timeout=60.0
                )
            except Exception as e2:
                return JSONResponse({"error": f"Fix failed: {e2}"}, status_code=503)

        result = clean_thinking(result)
        result = strip_code_fences(result)

        fix_plan = None
        try:
            fix_plan = json.loads(result)
        except:
            m = re.search(r'\{[\s\S]*\}', result)
            if m:
                try: fix_plan = json.loads(m.group(0))
                except: pass
        if not fix_plan or not fix_plan.get("files"):
            return JSONResponse({"error": "Could not parse fix response", "raw": result[:500]}, status_code=500)

        credits = await get_credits(user_id)
        return {"fix": fix_plan, "credits": {"remaining": credits["remaining"], "total": credits["total"], "tier": credits["tier"]}}
    except Exception as e:
        print(f"[Studio][Fix] Error: {e}")
        return JSONResponse({"error": str(e)}, status_code=500)

# ── CREDITS ENDPOINT ──────────────────────────────────────────────────────────
@router.get("/credits/{user_id}")
async def get_user_credits(user_id: str, email: str = ""):
    """Get credit status + Studio plan/features for a user."""
    credits = await get_credits(user_id, email=email)
    return {"user_id": user_id, **credits}

# ── USAGE ENDPOINT (token stats for display) ──────────────────────────────────
@router.get("/usage/{user_id}")
async def get_studio_usage(user_id: str, email: str = ""):
    """Server truth for badge + stats: credits remaining and tokens used this month."""
    credits = await get_credits(user_id, email=email)
    return {
        "user_id": user_id,
        "credits_remaining": credits["remaining"],
        "credits_total": credits["total"],
        "tokens_used": credits.get("tokens_used", 0),
        "tokens_monthly": credits.get("tokens_monthly", _tokens_budget(credits["tier"])),
        "tokens_input": credits.get("tokens_input", 0),
        "tokens_output": credits.get("tokens_output", 0),
        "tier": credits["tier"],
        "plan": credits.get("plan", credits["tier"]),
        "month": credits["month"],
        "features": credits.get("features") or _features_for_plan(credits["tier"]),
        "free_credits": PLAN_CREDITS["free"],
        "pro_credits": PLAN_CREDITS["pro"],
        "max_credits": PLAN_CREDITS["max"],
        "business_credits": PLAN_CREDITS["business"],
        "plan_credits": PLAN_CREDITS,
        "plan_tokens": PLAN_TOKENS,
    }

# ── STATS ENDPOINT (admin) ────────────────────────────────────────────────────
@router.get("/stats")
async def get_stats(user_id: str = ""):
    """Global studio statistics. Admin-only if ADMIN_UIDS is set."""
    if ADMIN_UIDS and user_id and user_id not in ADMIN_UIDS:
        return JSONResponse({"error": "Admin only"}, status_code=403)
    try:
        doc = await _db_find_one("studio_stats", {"_id": "global"})
        if doc:
            doc.pop("_id", None)
            return doc
        return {"total_projects": 0, "total_users": 0, "active_users_7d": 0, "projects_today": 0, "projects_this_month": 0}
    except Exception as e:
        return JSONResponse({"error": str(e)}, 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 v2] Build request — uid:{user_id} | prompt:{user_prompt[:80]}")

    async def stream():
        # Rate limit
        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": "Studio builds web apps only. Try a website, dashboard, or web tool."}); return

        # Credit + plan check (polar_subs_studio)
        credits_info = await get_credits(user_id)
        remaining = credits_info["remaining"]
        total = credits_info["total"]
        tier = credits_info["tier"]
        features = credits_info.get("features") or _features_for_plan(tier)
        if remaining <= 0:
            yield sse("error", {"message": f"No credits remaining this month. You get {total} credits/month ({tier} plan). Upgrade for more builds.", "credits": {"remaining": 0, "total": total, "tier": tier, "plan": tier, "features": features}}); return

        # File caps by plan
        if tier == "business":
            max_files = max(MAX_FILES_PRO, 20)
        elif tier in ("pro", "max"):
            max_files = MAX_FILES_PRO
        else:
            max_files = MAX_FILES_FREE

        yield sse("status", {"message": "Initialising Studio Engine..."})
        yield sse("target", {"target": "web", "previewable": True})
        yield sse("status", {"message": "Designing your app architecture..."})
        yield sse("credits", {
            "remaining": remaining, "total": total, "tier": tier, "plan": tier,
            "features": features,
            "tokens_used": credits_info.get("tokens_used", 0),
            "tokens_monthly": credits_info.get("tokens_monthly", _tokens_budget(tier)),
        })
        build_in_tok = 0
        build_out_tok = 0

        # PLANNER
        plan = None
        for plan_attempt in range(1, 3):
            try:
                plan_out = await asyncio.to_thread(call_openrouter_raw,
                    [{"role": "system", "content": PLANNER_SYSTEM}, {"role": "user", "content": user_prompt}],
                    BUILD_MODEL, 4000)
                u = plan_out.get("usage") or {}
                build_in_tok += int(u.get("input_tokens") or 0)
                build_out_tok += int(u.get("output_tokens") or 0)
                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}/2)..."})
                await asyncio.sleep(2)

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

        plan["files"] = plan["files"][:max_files]
        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')})"})

        # FILE GENERATION (sequential)
        files = {}
        for idx, file_spec in enumerate(plan["files"], start=1):
            path = file_spec.get("path")
            if not path: continue
            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)
            written_ctx = ""
            if files:
                written_ctx = "\n=== WRITTEN SO FAR ===\n"
                for wp, wc in files.items():
                    snippet = wc if len(wc) < 8000 else wc[:4000] + "\n... [TRUNCATED] ...\n" + wc[-4000:]
                    written_ctx += f"--- {wp} ---\n{snippet}\n"
                written_ctx += "===\n"

            user_msg = f"""USER PROMPT:\n{user_prompt}\n\nPROJECT PLAN:\n{plan_json}\n\nALL FILES:\n{json.dumps(all_paths)}\n{written_ctx}\nFILE TO GENERATE:\n  path: {path}\n  purpose: {file_spec.get('purpose', '')}\n  key_features: {file_spec.get('key_features', [])}\n\nOutput ONLY the complete raw content of `{path}`. No fences."""
            messages = [{"role": "system", "content": FILE_SYSTEM}, {"role": "user", "content": user_msg}]

            try:
                out = await call_with_retry(messages, target_model, MAX_TOKENS_PER_CALL)
                u = out.get("usage") or {}
                build_in_tok += int(u.get("input_tokens") or 0)
                build_out_tok += int(u.get("output_tokens") or 0)
                accumulated = clean_thinking(out["content"])
                finish = out["finish_reason"]
            except Exception as e:
                print(f"[Studio] Error on {path}: {e}")
                yield sse("status", {"message": f"Fallback on {path}..."})
                try:
                    out = await call_with_retry(messages, FALLBACK_MODEL, MAX_TOKENS_PER_CALL)
                    accumulated = clean_thinking(out["content"])
                    finish = out["finish_reason"]
                except:
                    yield sse("log", {"message": f"Skipped {path}"})
                    continue

            # Continuation
            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 `{path}`.\nLast chars (do NOT repeat):\n---\n{tail}\n---\nContinue the rest. Raw content only."}]
                await asyncio.sleep(1)
                try:
                    out = await call_with_retry(cont_messages, target_model, MAX_TOKENS_PER_CALL)
                    u = out.get("usage") or {}
                    build_in_tok += int(u.get("input_tokens") or 0)
                    build_out_tok += int(u.get("output_tokens") or 0)
                    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: 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)"})

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

        # CONSUME CREDIT (only on successful build) + persist tokens
        consumed = await consume_credit(user_id, "build")
        if not consumed:
            yield sse("error", {"message": "Could not consume credit."}); return
        await record_tokens(user_id, build_in_tok, build_out_tok)

        # WIRE + SAVE
        files = inject_file_references(files)
        yield sse("status", {"message": "Committing to storage..."})
        try:
            expires_at = datetime.utcnow() + timedelta(hours=24)
            project_id = await save_project_to_mongo(user_id, user_prompt, plan, files, expires_at, tier)
            await _update_stats()
            credits = await get_credits(user_id)
            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(),
                "tokens": {"input": build_in_tok, "output": build_out_tok, "total": build_in_tok + build_out_tok},
                "credits": {
                    "remaining": credits["remaining"],
                    "total": credits["total"],
                    "tier": credits["tier"],
                    "plan": credits.get("plan", credits["tier"]),
                    "tokens_used": credits.get("tokens_used", 0),
                    "tokens_monthly": credits.get("tokens_monthly", _tokens_budget(credits["tier"])),
                    "features": credits.get("features") or features,
                }})
        except Exception as e:
            print(f"[Studio] Save error: {e}")
            yield sse("done", {"message": "Build complete (cached)", "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_find_one("studio_projects", {"_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: 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: return JSONResponse({"error": "Error"}, 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_find_one("studio_projects", {"_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_delete_one("studio_projects", {"_id": ObjectId(project_id)})
        result = await _db_delete_many("studio_files", {"project_id": project_id})
        return {"deleted": True, "files_removed": result.deleted_count}
    except: return JSONResponse({"error": "Delete failed"}, status_code=500)

@router.get("/project/{project_id}/download")
async def download_project(project_id: str):
    import io, zipfile
    try:
        p = await _db_find_one("studio_projects", {"_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:
            files_raw = await _db_find("studio_files", {"project_id": project_id}, sort=("path", 1))
            for f in files_raw:
                path = (f.get("path") or "").lstrip("/").replace("\\", "/")
                if not path or ".." in path.split("/"): continue
                content = f.get("content", "")
                data = content.encode("utf-8") if isinstance(content, str) else content
                zf.writestr(path, data)
        buf.seek(0)
        safe_name = "".join(c if c.isalnum() or c in "-_" else "_" for c in (p.get("app_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: return JSONResponse({"error": "Download failed"}, status_code=500)