main.py

import json
import logging
import os
import re
import subprocess
import sys
import uuid
from datetime import datetime
from pathlib import Path
from logging.handlers import RotatingFileHandler
from typing import AsyncGenerator, Any, Optional
import shutil

from dotenv import load_dotenv
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from fastapi.staticfiles import StaticFiles
from openai import AsyncOpenAI
from pydantic import BaseModel
import uvicorn
import asyncio

load_dotenv()

# =============================================================================
# SCHEMA STORE - 本地存储管理
# =============================================================================
SCHEMA_STORE_DIR = Path(__file__).parent / "schema_store"
SCHEMA_STORE_DIR.mkdir(exist_ok=True)


class SchemaStore:
    """
    管理 UI Schema 的本地存储，支持分片存储和递归生成。

    增强功能:
    - $ref 引用机制: 骨架中使用引用指向分片
    - 上下文压缩: 为 LLM 提供精简的上下文摘要
    - 增量加载: 支持按需加载和合并分片
    - 依赖追踪: 追踪分片之间的依赖关系

    目录结构:
    schema_store/
    ├── {session_id}/
    │   ├── session.json      # 会话元数据
    │   ├── skeleton.json     # 主骨架结构 (含 $ref 引用)
    │   ├── skeleton_refs.json # 骨架引用版本 (用于增量加载)
    │   ├── fragments/        # 分片存储
    │   │   ├── {fragment_id}.json
    │   │   └── ...
    │   ├── context/          # 上下文摘要
    │   │   └── summary.json  # 压缩的上下文信息
    │   └── merged.json       # 最终合并结果
    """

    def __init__(self, session_id: str):
        self.session_id = session_id
        self.session_dir = SCHEMA_STORE_DIR / session_id
        self.fragments_dir = self.session_dir / "fragments"
        self.context_dir = self.session_dir / "context"
        self._ensure_dirs()
        self._fragment_index: dict[str, dict] = {}  # 内存中的分片索引

    def _ensure_dirs(self):
        """确保目录存在"""
        self.session_dir.mkdir(parents=True, exist_ok=True)
        self.fragments_dir.mkdir(exist_ok=True)
        self.context_dir.mkdir(exist_ok=True)

    def save_session_meta(self, meta: dict):
        """保存会话元数据"""
        meta["updated_at"] = datetime.now().isoformat()
        with open(self.session_dir / "session.json", "w", encoding="utf-8") as f:
            json.dump(meta, f, ensure_ascii=False, indent=2)

    def load_session_meta(self) -> dict | None:
        """加载会话元数据"""
        path = self.session_dir / "session.json"
        if path.exists():
            with open(path, "r", encoding="utf-8") as f:
                return json.load(f)
        return None

    def save_skeleton(self, skeleton: dict):
        """保存主骨架"""
        with open(self.session_dir / "skeleton.json", "w", encoding="utf-8") as f:
            json.dump(skeleton, f, ensure_ascii=False, indent=2)

    def load_skeleton(self) -> dict | None:
        """加载主骨架"""
        path = self.session_dir / "skeleton.json"
        if path.exists():
            with open(path, "r", encoding="utf-8") as f:
                return json.load(f)
        return None

    def save_skeleton_with_refs(self, skeleton: dict):
        """
        保存带有 $ref 引用的骨架版本
        将 $placeholder 转换为 $ref 引用，便于增量加载
        """
        skeleton_refs = self._convert_placeholders_to_refs(skeleton)
        with open(self.session_dir / "skeleton_refs.json", "w", encoding="utf-8") as f:
            json.dump(skeleton_refs, f, ensure_ascii=False, indent=2)
        return skeleton_refs

    def load_skeleton_with_refs(self) -> dict | None:
        """加载带引用的骨架"""
        path = self.session_dir / "skeleton_refs.json"
        if path.exists():
            with open(path, "r", encoding="utf-8") as f:
                return json.load(f)
        return None

    def _convert_placeholders_to_refs(self, schema: dict) -> dict:
        """将 $placeholder 转换为 $ref 引用"""
        if not isinstance(schema, dict):
            return schema

        result = {}
        for key, value in schema.items():
            if key == "$placeholder":
                # 转换为 $ref 引用
                result["$ref"] = f"#/fragments/{value}"
            elif key == "children" and isinstance(value, list):
                result[key] = [self._convert_placeholders_to_refs(child) if isinstance(child, dict) else child for child in value]
            elif isinstance(value, dict):
                result[key] = self._convert_placeholders_to_refs(value)
            else:
                result[key] = value
        return result

    def save_fragment(self, fragment_id: str, fragment: dict, meta: dict = None):
        """
        保存一个 schema 分片

        Args:
            fragment_id: 分片ID (通常是 placeholder_id)
            fragment: 分片内容
            meta: 分片元数据 (深度、父节点、依赖等)
        """
        # 生成分片摘要用于上下文
        summary = self._generate_fragment_summary(fragment)

        fragment_data = {
            "id": fragment_id,
            "content": fragment,
            "summary": summary,
            "meta": meta or {},
            "created_at": datetime.now().isoformat()
        }
        with open(self.fragments_dir / f"{fragment_id}.json", "w", encoding="utf-8") as f:
            json.dump(fragment_data, f, ensure_ascii=False, indent=2)

        # 更新内存索引
        self._fragment_index[fragment_id] = {
            "summary": summary,
            "meta": meta or {},
            "has_refs": self._has_nested_refs(fragment)
        }

        # 更新上下文摘要
        self._update_context_summary()

    def _generate_fragment_summary(self, fragment: dict, max_depth: int = 2) -> dict:
        """
        生成分片的结构摘要，用于上下文压缩
        只保留类型、关键属性和子组件类型
        """
        if not isinstance(fragment, dict):
            return {"type": "text", "preview": str(fragment)[:50]}

        summary = {
            "type": fragment.get("type", "unknown"),
            "props_keys": list(fragment.get("props", {}).keys())[:5],
        }

        children = fragment.get("children", [])
        if isinstance(children, str):
            summary["children_type"] = "text"
            summary["children_preview"] = children[:30] + "..." if len(children) > 30 else children
        elif isinstance(children, list):
            summary["children_type"] = "array"
            summary["children_count"] = len(children)
            if max_depth > 0:
                summary["children_types"] = [
                    self._generate_fragment_summary(c, max_depth - 1).get("type", "?")
                    if isinstance(c, dict) else "text"
                    for c in children[:5]
                ]

        # 检查是否有嵌套的 placeholder
        if self._has_nested_refs(fragment):
            summary["has_nested_refs"] = True

        return summary

    def _has_nested_refs(self, schema: dict) -> bool:
        """检查 schema 中是否有嵌套的 $placeholder 或 $ref"""
        if not isinstance(schema, dict):
            return False
        if "$placeholder" in schema or "$ref" in schema:
            return True
        for value in schema.values():
            if isinstance(value, dict) and self._has_nested_refs(value):
                return True
            elif isinstance(value, list):
                for item in value:
                    if isinstance(item, dict) and self._has_nested_refs(item):
                        return True
        return False

    def _update_context_summary(self):
        """更新上下文摘要文件"""
        summary = {
            "session_id": self.session_id,
            "updated_at": datetime.now().isoformat(),
            "fragment_count": len(self._fragment_index),
            "fragments": {}
        }

        for fid, info in self._fragment_index.items():
            summary["fragments"][fid] = {
                "summary": info["summary"],
                "depth": info["meta"].get("depth", 0),
                "has_refs": info["has_refs"]
            }

        with open(self.context_dir / "summary.json", "w", encoding="utf-8") as f:
            json.dump(summary, f, ensure_ascii=False, indent=2)

    def load_fragment(self, fragment_id: str) -> dict | None:
        """加载一个分片"""
        path = self.fragments_dir / f"{fragment_id}.json"
        if path.exists():
            with open(path, "r", encoding="utf-8") as f:
                return json.load(f)
        return None

    def list_fragments(self) -> list[str]:
        """列出所有分片ID"""
        return [f.stem for f in self.fragments_dir.glob("*.json")]

    def get_fragments_by_depth(self, depth: int) -> list[dict]:
        """获取指定深度的所有分片"""
        fragments = []
        for fid in self.list_fragments():
            frag = self.load_fragment(fid)
            if frag and frag.get("meta", {}).get("depth") == depth:
                fragments.append(frag)
        return fragments

    def save_merged(self, schema: dict):
        """保存最终合并的 schema"""
        with open(self.session_dir / "merged.json", "w", encoding="utf-8") as f:
            json.dump(schema, f, ensure_ascii=False, indent=2)

    def load_merged(self) -> dict | None:
        """加载合并后的 schema"""
        path = self.session_dir / "merged.json"
        if path.exists():
            with open(path, "r", encoding="utf-8") as f:
                return json.load(f)
        return None

    def resolve_refs(self, schema: dict) -> dict:
        """
        解析 schema 中的 $ref 引用，返回完整的 schema
        支持递归解析嵌套引用
        """
        if not isinstance(schema, dict):
            return schema

        # 如果是 $ref 引用，加载对应分片
        if "$ref" in schema:
            ref_path = schema["$ref"]
            if ref_path.startswith("#/fragments/"):
                fragment_id = ref_path.replace("#/fragments/", "")
                frag = self.load_fragment(fragment_id)
                if frag and frag.get("content"):
                    # 递归解析分片中的引用
                    return self.resolve_refs(frag["content"])
            return schema

        # 递归处理子节点
        result = {}
        for key, value in schema.items():
            if key == "children" and isinstance(value, list):
                result[key] = [self.resolve_refs(child) if isinstance(child, dict) else child for child in value]
            elif isinstance(value, dict):
                result[key] = self.resolve_refs(value)
            else:
                result[key] = value
        return result

    def get_context_summary(self, max_fragments: int = 5) -> str:
        """
        获取上下文摘要，用于传递给 LLM
        只包含最近的几个分片的类型和结构，不包含完整内容
        """
        fragments = []
        for fid in self.list_fragments()[-max_fragments:]:
            frag = self.load_fragment(fid)
            if frag and frag.get("content"):
                content = frag["content"]
                # 只提取类型和顶层结构
                summary = self._summarize_schema(content, max_depth=2)
                fragments.append(f"- {fid}: {summary}")

        return "\n".join(fragments) if fragments else "No fragments yet"

    def get_compressed_context(self, max_tokens: int = 2000) -> dict:
        """
        获取压缩的上下文信息，用于传递给 LLM
        包含骨架结构摘要和已完成分片的类型信息

        Args:
            max_tokens: 估算的最大 token 数

        Returns:
            压缩的上下文字典
        """
        context = {
            "session_id": self.session_id,
            "skeleton_structure": None,
            "completed_fragments": [],
            "pending_refs": []
        }

        # 加载骨架摘要
        skeleton = self.load_skeleton()
        if skeleton:
            context["skeleton_structure"] = self._summarize_schema(skeleton, max_depth=3)

        # 加载分片摘要
        context_file = self.context_dir / "summary.json"
        if context_file.exists():
            with open(context_file, "r", encoding="utf-8") as f:
                summary_data = json.load(f)
                for fid, info in summary_data.get("fragments", {}).items():
                    context["completed_fragments"].append({
                        "id": fid,
                        "type": info["summary"].get("type"),
                        "depth": info.get("depth", 0)
                    })
                    if info.get("has_refs"):
                        context["pending_refs"].append(fid)

        return context

    def _summarize_schema(self, schema: dict, depth: int = 0, max_depth: int = 2) -> str:
        """生成 schema 的简短摘要"""
        if depth > max_depth or not isinstance(schema, dict):
            return "..."

        comp_type = schema.get("type", "unknown")
        children = schema.get("children", [])

        if isinstance(children, str):
            return f"{comp_type}(\"{children[:20]}...\")" if len(children) > 20 else f"{comp_type}(\"{children}\")"
        elif isinstance(children, list) and children:
            child_types = [c.get("type", "?") if isinstance(c, dict) else "text" for c in children[:3]]
            suffix = f"+{len(children)-3}" if len(children) > 3 else ""
            return f"{comp_type}[{', '.join(child_types)}{suffix}]"
        else:
            return comp_type

    @classmethod
    def list_sessions(cls) -> list[dict]:
        """列出所有会话"""
        sessions = []
        for session_dir in SCHEMA_STORE_DIR.iterdir():
            if session_dir.is_dir():
                meta_path = session_dir / "session.json"
                if meta_path.exists():
                    with open(meta_path, "r", encoding="utf-8") as f:
                        meta = json.load(f)
                        meta["session_id"] = session_dir.name
                        sessions.append(meta)
        return sorted(sessions, key=lambda x: x.get("updated_at", ""), reverse=True)

    @classmethod
    def delete_session(cls, session_id: str):
        """删除会话"""
        session_dir = SCHEMA_STORE_DIR / session_id
        if session_dir.exists():
            shutil.rmtree(session_dir)

# =============================================================================
# LOGGING CONFIGURATION
# =============================================================================
LOG_DIR = Path(__file__).parent / "logs"
LOG_DIR.mkdir(exist_ok=True)

# Get log level from environment variable (default: INFO for console, DEBUG for file)
# Options: DEBUG, INFO, WARNING, ERROR
LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO").upper()
LOG_LEVEL_NUM = getattr(logging, LOG_LEVEL, logging.INFO)

# Create formatters
detailed_formatter = logging.Formatter(
    '%(asctime)s | %(levelname)-8s | %(name)-20s | %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
console_formatter = logging.Formatter(
    '%(asctime)s | %(levelname)-8s | %(message)s',
    datefmt='%H:%M:%S'
)

# Root logger setup
root_logger = logging.getLogger()
root_logger.setLevel(logging.DEBUG)

# Console handler (level from env var)
console_handler = logging.StreamHandler(sys.stdout)
console_handler.setLevel(LOG_LEVEL_NUM)
console_handler.setFormatter(console_formatter)
root_logger.addHandler(console_handler)

# File handler for all logs (DEBUG level, rotating)
all_log_file = LOG_DIR / "realtimeui.log"
file_handler = RotatingFileHandler(
    all_log_file,
    maxBytes=10 * 1024 * 1024,  # 10MB
    backupCount=5,
    encoding='utf-8'
)
file_handler.setLevel(logging.DEBUG)
file_handler.setFormatter(detailed_formatter)
root_logger.addHandler(file_handler)

# Separate file for LLM interactions (for analysis)
llm_log_file = LOG_DIR / "llm_calls.log"
llm_handler = RotatingFileHandler(
    llm_log_file,
    maxBytes=50 * 1024 * 1024,  # 50MB (LLM logs can be large)
    backupCount=10,
    encoding='utf-8'
)
llm_handler.setLevel(logging.DEBUG)
llm_handler.setFormatter(detailed_formatter)

# Create loggers
logger = logging.getLogger("realtimeui")
llm_logger = logging.getLogger("realtimeui.llm")
llm_logger.addHandler(llm_handler)
agent_logger = logging.getLogger("realtimeui.agent")

# Reduce noise from third-party libraries
logging.getLogger("httpx").setLevel(logging.WARNING)
logging.getLogger("httpcore").setLevel(logging.WARNING)
logging.getLogger("openai").setLevel(logging.WARNING)
logging.getLogger("uvicorn.access").setLevel(logging.WARNING)

logger.info(f"Log level: {LOG_LEVEL} (console), DEBUG (file)")
logger.info(f"Log files: {LOG_DIR}")

app = FastAPI(title="RealtimeUI API", version="0.1.0")

# Check if running in production mode (serving static files)
FRONTEND_DIR = Path(__file__).parent / "frontend"
DIST_DIR = FRONTEND_DIR / "dist"
IS_PRODUCTION = DIST_DIR.exists() and (DIST_DIR / "index.html").exists()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["http://localhost:5173", "http://localhost:8000"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

client = AsyncOpenAI(
    api_key=os.getenv("OPENAI_API_KEY"),
    base_url=os.getenv("OPENAI_BASE_URL"),
)

MODEL = os.getenv("OPENAI_MODEL", "qwen3-max")

# Support multiple models for parallel agents
# Format: comma-separated list of models, e.g., "qwen3-max,qwen-plus,qwen-turbo"
EXPANSION_MODELS = os.getenv("EXPANSION_MODELS", "").split(",") if os.getenv("EXPANSION_MODELS") else []
EXPANSION_MODELS = [m.strip() for m in EXPANSION_MODELS if m.strip()]

# If no expansion models specified, use the main model
if not EXPANSION_MODELS:
    EXPANSION_MODELS = [MODEL]

logger.info(f"Main model: {MODEL}")
logger.info(f"Expansion models: {EXPANSION_MODELS}")

COMPONENTS_DIR = Path(__file__).parent / "components"

# =============================================================================
# HIERARCHICAL UI GENERATION (分层生成)
# =============================================================================
# Multi-Agent Architecture:
#
# 1. SkeletonAgent (架构 Agent)
#    - Generates high-level structure with $placeholder markers
#    - Only creates the skeleton, no actual content
#    - Fast, small context window usage
#
# 2. ExpansionAgent (展开 Agent) - Multiple instances run in parallel
#    - Each agent expands one placeholder
#    - Isolated context per agent
#    - Streams results back to orchestrator
#
# 3. Orchestrator (协调器)
#    - Manages agent lifecycle
#    - Collects results via event queue
#    - Merges expanded content into final schema
#
# Benefits:
#    - True parallel execution
#    - Context isolation per agent
#    - Progressive rendering
#    - Failure isolation
# =============================================================================

import asyncio
from asyncio import Queue as AsyncQueue
from dataclasses import dataclass
from typing import AsyncGenerator, Any
from abc import ABC, abstractmethod


@dataclass
class AgentEvent:
    """Event emitted by an agent."""
    agent_id: str
    event_type: str
    data: dict


class BaseAgent(ABC):
    """Base class for all agents."""

    def __init__(self, agent_id: str, client: AsyncOpenAI, model: str):
        self.agent_id = agent_id
        self.client = client
        self.model = model

    @abstractmethod
    async def run(self) -> AsyncGenerator[AgentEvent, None]:
        """Run the agent and yield events."""
        pass


class SkeletonAgent(BaseAgent):
    """Agent that generates the UI skeleton structure with placeholders."""

    SYSTEM_PROMPT = """You are a UI architect agent. Generate a HIGH-LEVEL skeleton structure for complex UIs.

## CRITICAL: Close-First-Then-Fill Principle (先闭合再填充)
You MUST follow this generation order:
1. First output the complete CLOSED structure (all brackets matched)
2. Use $placeholder markers for content that will be filled later
3. NEVER leave brackets unclosed - always complete the structure first

## Your Role
You are the ARCHITECTURE agent. Your job is to:
1. Understand the overall UI requirements
2. Design the high-level structure with COMPLETE, CLOSED JSON
3. Mark complex sections with $placeholder for other agents to expand

## Output Format
Return a COMPLETE, CLOSED JSON object with $placeholder markers:
{{
  "type": "ComponentName",
  "props": {{ }},
  "children": [
    {{
      "$placeholder": "unique_id",
      "$description": "Detailed description for the expansion agent"
    }}
  ]
}}

## Generation Order (IMPORTANT!)
1. First write the opening braces and type
2. Then write props (can be empty {{}})
3. Then write children array with placeholders
4. CLOSE all brackets before moving to next sibling
5. Each placeholder should be a complete, closed object

## Rules
1. Use $placeholder for sections that need detailed content (3+ components)
2. Simple elements (single Text, Button, etc.) can be inline
3. Each $placeholder must have a unique id and clear $description
4. $description should be detailed enough for another agent to implement
5. Keep the skeleton shallow (max 2-3 levels before placeholder)
6. IMPORTANT: For sidebar navigation, ALWAYS use Tabs with orientation="vertical"

## Available Layout Components
- Row: Horizontal container. Props: gap ("sm"|"md"|"lg")
- Col: Vertical container. Props: gap ("sm"|"md"|"lg")
- Card, CardHeader, CardContent, CardFooter: Card sections
- Tabs: Tab container with switchable content. Props: defaultValue (required), orientation ("horizontal"|"vertical")
- TabsList: Container for tab triggers
- TabsTrigger: Clickable tab button. Props: value (required)
- TabsContent: Content panel. Props: value (required)

## CRITICAL: Sidebar Navigation Pattern
For ANY UI with sidebar navigation, you MUST use Tabs with vertical orientation:
{{
  "type": "Tabs",
  "props": {{"defaultValue": "first_tab", "orientation": "vertical"}},
  "children": [
    {{"type": "TabsList", "children": [
      {{"type": "TabsTrigger", "props": {{"value": "tab1"}}, "children": "Tab 1"}},
      {{"type": "TabsTrigger", "props": {{"value": "tab2"}}, "children": "Tab 2"}}
    ]}},
    {{"type": "TabsContent", "props": {{"value": "tab1"}}, "children": [{{"$placeholder": "...", "$description": "..."}}]}},
    {{"type": "TabsContent", "props": {{"value": "tab2"}}, "children": [{{"$placeholder": "...", "$description": "..."}}]}}
  ]
}}

## Example
User: "Create a dashboard with sidebar navigation, stats cards, and data table"
{{
  "type": "Tabs",
  "props": {{"defaultValue": "overview", "orientation": "vertical"}},
  "children": [
    {{"type": "TabsList", "children": [
      {{"type": "TabsTrigger", "props": {{"value": "overview"}}, "children": "Overview"}},
      {{"type": "TabsTrigger", "props": {{"value": "orders"}}, "children": "Orders"}},
      {{"type": "TabsTrigger", "props": {{"value": "products"}}, "children": "Products"}}
    ]}},
    {{"type": "TabsContent", "props": {{"value": "overview"}}, "children": [
      {{
        "$placeholder": "stats_section",
        "$description": "Statistics cards showing key metrics: total revenue, total orders, total customers, conversion rate. Each card should have title, value, trend indicator."
      }},
      {{
        "$placeholder": "recent_orders_table",
        "$description": "Table showing recent orders with columns: Order ID, Customer, Status (badge), Amount, Date. Show 5 sample rows."
      }}
    ]}},
    {{"type": "TabsContent", "props": {{"value": "orders"}}, "children": [
      {{
        "$placeholder": "orders_content",
        "$description": "Full orders management view with search input, filter dropdown, and orders table."
      }}
    ]}},
    {{"type": "TabsContent", "props": {{"value": "products"}}, "children": [
      {{
        "$placeholder": "products_content",
        "$description": "Products management view with product grid or table."
      }}
    ]}}
  ]
}}

Now generate the skeleton for the user's request. Output ONLY the JSON."""

    def __init__(self, agent_id: str, client: AsyncOpenAI, model: str, prompt: str):
        super().__init__(agent_id, client, model)
        self.prompt = prompt

    async def run(self) -> AsyncGenerator[AgentEvent, None]:
        """Generate skeleton structure with streaming."""
        import time
        start_time = time.time()
        agent_logger.info(f"SkeletonAgent STARTED | prompt_length={len(self.prompt)}")

        yield AgentEvent(self.agent_id, "start", {"message": "SkeletonAgent starting..."})

        messages = [
            {"role": "system", "content": self.SYSTEM_PROMPT},
            {"role": "user", "content": self.prompt}
        ]

        # Log full prompt to LLM
        llm_logger.info(f"LLM_REQUEST | agent={self.agent_id} | model={self.model}")
        llm_logger.debug(f"LLM_SYSTEM_PROMPT | agent={self.agent_id} | length={len(self.SYSTEM_PROMPT)}")
        llm_logger.debug(f"LLM_SYSTEM_PROMPT_CONTENT | agent={self.agent_id} |\n{self.SYSTEM_PROMPT[:2000]}...")
        llm_logger.info(f"LLM_USER_PROMPT | agent={self.agent_id} | prompt={self.prompt}")

        try:
            stream = await self.client.chat.completions.create(
                model=self.model,
                messages=messages,
                temperature=0.7,
                stream=True,
            )

            full_content = ""
            token_count = 0
            async for chunk in stream:
                if chunk.choices and chunk.choices[0].delta.content:
                    token = chunk.choices[0].delta.content
                    full_content += token
                    token_count += 1

                    # Log every 50 tokens for progress tracking
                    if token_count % 50 == 0:
                        elapsed = time.time() - start_time
                        agent_logger.debug(f"SkeletonAgent: {token_count} tokens ({elapsed:.1f}s)")

                    # Try to parse partial JSON
                    parsed = parse_partial_json(full_content)
                    yield AgentEvent(self.agent_id, "token", {
                        "token": token,
                        "full": full_content,
                        "parsed": parsed
                    })

            # Final parse
            skeleton = parse_partial_json(full_content)
            elapsed = time.time() - start_time

            # Log complete LLM response
            agent_logger.info(f"SkeletonAgent COMPLETED | tokens={token_count} | time={elapsed:.1f}s")
            llm_logger.info(f"LLM_RESPONSE | agent={self.agent_id} | tokens={token_count} | time={elapsed:.1f}s")
            llm_logger.info(f"LLM_OUTPUT_CONTENT | agent={self.agent_id} |\n{full_content}")

            yield AgentEvent(self.agent_id, "complete", {
                "skeleton": skeleton,
                "raw": full_content
            })

        except Exception as e:
            import traceback
            agent_logger.error(f"SkeletonAgent ERROR: {e}")
            agent_logger.error(traceback.format_exc())
            llm_logger.error(f"LLM_ERROR | agent={self.agent_id} | error={e}")
            yield AgentEvent(self.agent_id, "error", {"error": str(e)})


def create_openai_client():
    """Create a new AsyncOpenAI client instance."""
    return AsyncOpenAI(
        api_key=os.getenv("OPENAI_API_KEY"),
        base_url=os.getenv("OPENAI_BASE_URL"),
    )


class ExpansionAgent(BaseAgent):
    """Agent that expands a single placeholder into detailed UI components."""

    SYSTEM_PROMPT = """You are a UI component expansion agent. Your job is to expand a placeholder into detailed UI components.

## CRITICAL: Close-First-Then-Fill Principle (先闭合再填充)
You MUST follow this generation order:
1. First output the COMPLETE CLOSED structure (all brackets matched)
2. Fill in content values AFTER the structure is complete
3. NEVER leave brackets unclosed - always complete each object/array before starting the next

## Generation Order (IMPORTANT!)
When generating JSON, follow this exact order:
1. Open the root object {{
2. Write "type": "..."
3. Write "props": {{...}} (complete and close it)
4. Write "children": [ and for each child:
   a. Open child {{
   b. Write all properties
   c. Close child }}
   d. Add comma if not last
5. Close children ]
6. Close root }}

## Your Role
You are an EXPANSION agent. Another agent has created the overall UI structure and marked a section for you to implement.

## Context
- Overall UI purpose: {overall_purpose}
- Section to expand: {placeholder_id}
- Section description: {placeholder_description}
- Current depth: {current_depth} / Max depth: {max_depth}
- Context summary: {context_summary}

## Output Format
Return ONLY a valid, COMPLETE, CLOSED JSON object (the expanded UI section):
{{
  "type": "ComponentName",
  "props": {{ }},
  "children": [ ]
}}

## Recursive Expansion (IMPORTANT!)
If a section is too complex to fully implement (e.g., a complete sub-page, complex form, detailed dashboard section),
you can use $placeholder markers for further expansion:
{{
  "$placeholder": "unique_sub_id",
  "$description": "Detailed description of what this sub-section should contain"
}}

Rules for using $placeholder:
1. Only use if the section would require 10+ components to implement fully
2. Give each placeholder a unique, descriptive ID
3. Provide detailed $description for the next agent
4. Current depth is {current_depth}, max is {max_depth}. If depth >= max_depth, do NOT use placeholders.

## Available Components

### Layout
- Col: Vertical container. Props: gap ("sm"|"md"|"lg"), align, justify
- Row: Horizontal container. Props: gap, align, justify, wrap
- Card, CardHeader, CardContent, CardFooter, CardTitle, CardDescription

### Navigation
- Tabs: Tab container. Props: defaultValue (required), orientation ("horizontal"|"vertical")
- TabsList: Container for tab triggers
- TabsTrigger: Tab button. Props: value (required)
- TabsContent: Tab panel. Props: value (required)

### Basic
- Text: Props: as ("p"|"h1"|"h2"|"h3"), variant ("default"|"muted"), size ("sm"|"base"|"lg"|"xl"|"2xl")
- Button: Props: variant ("default"|"outline"|"secondary"|"destructive"|"ghost"), size ("sm"|"default"|"lg")
- Badge: Props: variant ("default"|"secondary"|"destructive"|"success"|"warning")
- Avatar: Props: src, alt, fallback, size ("sm"|"default"|"lg")
- Separator: Props: orientation ("horizontal"|"vertical")

### Form
- Input: Props: placeholder, type ("text"|"email"|"password"|"number")
- Textarea: Props: placeholder, rows
- Select, SelectTrigger, SelectValue, SelectContent, SelectItem
- Checkbox: Props: label, description
- Switch: Props: label, description
- Label: Form label
- Slider: Props: min, max, step

### Feedback
- Alert: Props: variant ("default"|"info"|"success"|"warning"|"destructive"), title
- Progress: Props: value (0-100), showValue, label

### Data
- Table, TableHeader, TableBody, TableRow, TableHead, TableCell
- Statistic: Props: title, value, prefix, suffix, trend ("up"|"down"|"neutral"), trendValue, description

### Charts
- LineChart, BarChart, PieChart, AreaChart: Props: data (array), xKey, yKey, title

Now expand the placeholder into detailed components. Output ONLY the JSON."""

    def __init__(self, agent_id: str, client: AsyncOpenAI, model: str,
                 placeholder: dict, overall_purpose: str,
                 current_depth: int = 1, max_depth: int = 3,
                 context_summary: str = ""):
        super().__init__(agent_id, client, model)
        self.placeholder = placeholder
        self.overall_purpose = overall_purpose
        self.current_depth = current_depth
        self.max_depth = max_depth
        self.context_summary = context_summary
        # Create a dedicated client for this agent
        self.dedicated_client = create_openai_client()

    async def run(self) -> AsyncGenerator[AgentEvent, None]:
        """Expand placeholder with streaming."""
        import time
        placeholder_id = self.placeholder["id"]
        start_time = time.time()
        agent_logger.info(f"ExpansionAgent STARTED: {placeholder_id} (depth={self.current_depth}/{self.max_depth})")

        yield AgentEvent(self.agent_id, "start", {
            "placeholder_id": placeholder_id,
            "depth": self.current_depth,
            "message": f"ExpansionAgent starting for {placeholder_id} (depth {self.current_depth})..."
        })

        prompt = self.SYSTEM_PROMPT.format(
            overall_purpose=self.overall_purpose,
            placeholder_id=placeholder_id,
            placeholder_description=self.placeholder["description"],
            current_depth=self.current_depth,
            max_depth=self.max_depth,
            context_summary=self.context_summary or "No previous context"
        )

        messages = [{"role": "user", "content": prompt}]

        # Log full LLM request including prompt content
        llm_logger.info(f"LLM_REQUEST | agent={self.agent_id} | model={self.model} | placeholder={placeholder_id}")
        llm_logger.debug(f"LLM_PROMPT_LENGTH | agent={self.agent_id} | length={len(prompt)}")
        llm_logger.info(f"LLM_PROMPT_CONTENT | agent={self.agent_id} | placeholder={placeholder_id} |\n{prompt}")

        try:
            agent_logger.debug(f"{placeholder_id}: Calling LLM API with model {self.model}")
            stream = await self.dedicated_client.chat.completions.create(
                model=self.model,
                messages=messages,
                temperature=0.7,
                stream=True,
            )
            agent_logger.debug(f"{placeholder_id}: LLM stream created successfully")

            full_content = ""
            token_count = 0
            async for chunk in stream:
                if chunk.choices and chunk.choices[0].delta.content:
                    token = chunk.choices[0].delta.content
                    full_content += token
                    token_count += 1

                    # Log every 50 tokens
                    if token_count % 50 == 0:
                        elapsed = time.time() - start_time
                        agent_logger.debug(f"{placeholder_id}: {token_count} tokens ({elapsed:.1f}s)")

                    # Try to parse partial JSON
                    parsed = parse_partial_json(full_content)
                    yield AgentEvent(self.agent_id, "token", {
                        "placeholder_id": placeholder_id,
                        "token": token,
                        "full": full_content,
                        "parsed": parsed
                    })

            # Final parse
            expanded = parse_partial_json(full_content)
            elapsed = time.time() - start_time
            agent_logger.info(f"ExpansionAgent COMPLETED: {placeholder_id} | tokens={token_count} | time={elapsed:.1f}s")

            # Log complete LLM response content
            llm_logger.info(f"LLM_RESPONSE | agent={self.agent_id} | placeholder={placeholder_id} | tokens={token_count} | time={elapsed:.1f}s")
            llm_logger.info(f"LLM_OUTPUT_CONTENT | agent={self.agent_id} | placeholder={placeholder_id} |\n{full_content}")

            yield AgentEvent(self.agent_id, "complete", {
                "placeholder_id": placeholder_id,
                "expanded": expanded,
                "path": self.placeholder["path"],
                "raw": full_content
            })

        except Exception as e:
            import traceback
            agent_logger.error(f"ExpansionAgent ERROR: {placeholder_id} - {e}")
            agent_logger.error(traceback.format_exc())
            llm_logger.error(f"LLM_ERROR | agent={self.agent_id} | placeholder={placeholder_id} | error={e}")
            yield AgentEvent(self.agent_id, "error", {
                "placeholder_id": placeholder_id,
                "error": str(e)
            })


class DataBindingAgent(BaseAgent):
    """Agent that binds data to UI components (Select options, Table data, Chart data, etc.)."""

    SYSTEM_PROMPT = """You are a Data Binding agent. Your job is to analyze a UI schema and add realistic mock data to components that need it.

## Your Role
You receive a complete UI schema and must:
1. Find all components that need data (Select, Table, Chart, RadioGroup, etc.)
2. Generate realistic mock data based on the UI context
3. Return JSON Patch operations to add the data

## Components That Need Data

### Select / SelectContent
- Add SelectItem children with realistic options
- Example: For a "Status" filter, add options like "All", "Active", "Pending", "Completed"

### Table / TableBody
- Add TableRow children with realistic data
- Match the columns defined in TableHeader
- Generate 5-10 rows of sample data

### Charts (LineChart, BarChart, PieChart, AreaChart)
- Add "data" prop with array of data points
- Include appropriate keys (xKey, yKey values)

### RadioGroup
- Add RadioGroupItem children with value and label

### Checkbox groups
- Set appropriate default checked states

## Context
- Overall UI purpose: {overall_purpose}
- Current schema structure summary: {schema_summary}

## Output Format
Return ONLY a JSON array of patch operations:
[
  {{"op": "add", "path": "/children/0/children/1/props/data", "value": [...]}},
  {{"op": "replace", "path": "/children/2/children", "value": [...]}}
]

## Rules
1. Generate realistic data that matches the UI context (e.g., e-commerce data for e-commerce UI)
2. Use Chinese text for Chinese UIs, English for English UIs
3. Include 5-10 items for lists/tables
4. For charts, generate 6-12 data points
5. Make data look realistic (varied values, proper formats)

Now analyze the schema and generate data binding patches. Output ONLY the JSON array:"""

    def __init__(self, agent_id: str, client: AsyncOpenAI, model: str,
                 schema: dict, overall_purpose: str):
        super().__init__(agent_id, client, model)
        self.schema = schema
        self.overall_purpose = overall_purpose
        self.dedicated_client = create_openai_client()

    def _summarize_schema(self, schema: dict, depth: int = 0, max_depth: int = 3) -> str:
        """Create a summary of the schema structure for the prompt."""
        if depth > max_depth:
            return "..."

        result = []
        if isinstance(schema, dict):
            comp_type = schema.get("type", "unknown")
            props = schema.get("props", {})
            props_str = ", ".join(f"{k}={v}" for k, v in list(props.items())[:3])
            result.append(f"{'  ' * depth}{comp_type}({props_str})")

            children = schema.get("children", [])
            if isinstance(children, list):
                for child in children[:5]:  # Limit children shown
                    if isinstance(child, dict):
                        result.append(self._summarize_schema(child, depth + 1, max_depth))
                if len(children) > 5:
                    result.append(f"{'  ' * (depth + 1)}... and {len(children) - 5} more")
            elif isinstance(children, str):
                result.append(f"{'  ' * (depth + 1)}\"{children[:30]}...\"" if len(children) > 30 else f"{'  ' * (depth + 1)}\"{children}\"")

        return "\n".join(result)

    async def run(self) -> AsyncGenerator[AgentEvent, None]:
        """Analyze schema and generate data bindings."""
        import time
        start_time = time.time()
        agent_logger.info(f"DataBindingAgent STARTED")

        yield AgentEvent(self.agent_id, "start", {
            "message": "DataBindingAgent analyzing schema for data binding..."
        })

        schema_summary = self._summarize_schema(self.schema)
        prompt = self.SYSTEM_PROMPT.format(
            overall_purpose=self.overall_purpose,
            schema_summary=schema_summary
        )

        # Also include the full schema for reference
        full_prompt = f"{prompt}\n\n## Full Schema:\n```json\n{json.dumps(self.schema, ensure_ascii=False, indent=2)}\n```"

        messages = [{"role": "user", "content": full_prompt}]

        llm_logger.info(f"LLM_REQUEST | agent={self.agent_id} | model={self.model}")
        llm_logger.debug(f"LLM_PROMPT_LENGTH | agent={self.agent_id} | length={len(full_prompt)}")

        try:
            stream = await self.dedicated_client.chat.completions.create(
                model=self.model,
                messages=messages,
                temperature=0.7,
                stream=True,
            )

            full_content = ""
            token_count = 0
            async for chunk in stream:
                if chunk.choices and chunk.choices[0].delta.content:
                    token = chunk.choices[0].delta.content
                    full_content += token
                    token_count += 1

                    yield AgentEvent(self.agent_id, "token", {
                        "token": token,
                        "full": full_content
                    })

            elapsed = time.time() - start_time
            agent_logger.info(f"DataBindingAgent COMPLETED | tokens={token_count} | time={elapsed:.1f}s")
            llm_logger.info(f"LLM_RESPONSE | agent={self.agent_id} | tokens={token_count} | time={elapsed:.1f}s")
            llm_logger.info(f"LLM_OUTPUT_CONTENT | agent={self.agent_id} |\n{full_content}")

            # Parse patches
            patches = parse_partial_patches(full_content)

            yield AgentEvent(self.agent_id, "complete", {
                "patches": patches or [],
                "raw": full_content
            })

        except Exception as e:
            import traceback
            agent_logger.error(f"DataBindingAgent ERROR: {e}")
            agent_logger.error(traceback.format_exc())
            yield AgentEvent(self.agent_id, "error", {"error": str(e)})


class UXEnhancementAgent(BaseAgent):
    """Agent that enhances UX with interactions, states, and feedback."""

    SYSTEM_PROMPT = """You are a UX Enhancement agent. Your job is to improve the user experience of a UI schema.

## Your Role
You receive a complete UI schema and must:
1. Add loading states where appropriate
2. Add empty states for lists/tables
3. Add form validation hints
4. Add hover/focus states
5. Add action feedback (button states, etc.)
6. Improve accessibility

## Enhancements to Add

### Buttons
- Add "loading" prop for action buttons
- Add appropriate icons
- Add tooltips via "title" prop

### Forms
- Add validation props (required, pattern, minLength, maxLength)
- Add helper text via description
- Add error states

### Tables/Lists
- Add empty state message
- Add loading skeleton

### Feedback
- Add success/error alerts after actions
- Add confirmation dialogs for destructive actions

### Accessibility
- Add aria-label where needed
- Ensure proper heading hierarchy

## Context
- Overall UI purpose: {overall_purpose}
- Current schema has data bindings applied

## Output Format
Return ONLY a JSON array of patch operations:
[
  {{"op": "add", "path": "/children/0/props/loading", "value": false}},
  {{"op": "add", "path": "/children/1/props/required", "value": true}}
]

## Rules
1. Don't break existing functionality
2. Add subtle enhancements, not major changes
3. Focus on common UX patterns
4. Keep it realistic and practical

Now analyze the schema and generate UX enhancement patches. Output ONLY the JSON array:"""

    def __init__(self, agent_id: str, client: AsyncOpenAI, model: str,
                 schema: dict, overall_purpose: str):
        super().__init__(agent_id, client, model)
        self.schema = schema
        self.overall_purpose = overall_purpose
        self.dedicated_client = create_openai_client()

    async def run(self) -> AsyncGenerator[AgentEvent, None]:
        """Analyze schema and generate UX enhancements."""
        import time
        start_time = time.time()
        agent_logger.info(f"UXEnhancementAgent STARTED")

        yield AgentEvent(self.agent_id, "start", {
            "message": "UXEnhancementAgent analyzing schema for UX improvements..."
        })

        prompt = self.SYSTEM_PROMPT.format(
            overall_purpose=self.overall_purpose
        )

        full_prompt = f"{prompt}\n\n## Current Schema:\n```json\n{json.dumps(self.schema, ensure_ascii=False, indent=2)}\n```"

        messages = [{"role": "user", "content": full_prompt}]

        llm_logger.info(f"LLM_REQUEST | agent={self.agent_id} | model={self.model}")
        llm_logger.debug(f"LLM_PROMPT_LENGTH | agent={self.agent_id} | length={len(full_prompt)}")

        try:
            stream = await self.dedicated_client.chat.completions.create(
                model=self.model,
                messages=messages,
                temperature=0.7,
                stream=True,
            )

            full_content = ""
            token_count = 0
            async for chunk in stream:
                if chunk.choices and chunk.choices[0].delta.content:
                    token = chunk.choices[0].delta.content
                    full_content += token
                    token_count += 1

                    yield AgentEvent(self.agent_id, "token", {
                        "token": token,
                        "full": full_content
                    })

            elapsed = time.time() - start_time
            agent_logger.info(f"UXEnhancementAgent COMPLETED | tokens={token_count} | time={elapsed:.1f}s")
            llm_logger.info(f"LLM_RESPONSE | agent={self.agent_id} | tokens={token_count} | time={elapsed:.1f}s")
            llm_logger.info(f"LLM_OUTPUT_CONTENT | agent={self.agent_id} |\n{full_content}")

            # Parse patches
            patches = parse_partial_patches(full_content)

            yield AgentEvent(self.agent_id, "complete", {
                "patches": patches or [],
                "raw": full_content
            })

        except Exception as e:
            import traceback
            agent_logger.error(f"UXEnhancementAgent ERROR: {e}")
            agent_logger.error(traceback.format_exc())
            yield AgentEvent(self.agent_id, "error", {"error": str(e)})


class AgentOrchestrator:
    """
    Orchestrates multiple agents for hierarchical UI generation with recursive expansion.

    Supports:
    - Recursive placeholder expansion (depth-first)
    - Local storage for schema fragments
    - Configurable max depth
    - Optional DataBinding and UXEnhancement phases
    """

    def __init__(
        self,
        client: AsyncOpenAI,
        model: str,
        max_depth: int = 3,
        enable_data_binding: bool = True,
        enable_ux_enhancement: bool = True
    ):
        self.client = client
        self.model = model
        self.max_depth = max_depth
        self.enable_data_binding = enable_data_binding
        self.enable_ux_enhancement = enable_ux_enhancement
        self.event_queue: AsyncQueue[AgentEvent] = AsyncQueue()
        self.expanded_results = {}
        self.store: SchemaStore | None = None
        self.total_expansions = 0
        self.current_depth_expansions = {}  # Track expansions per depth level

    async def _run_agent_and_queue_events(self, agent: BaseAgent):
        """Run an agent and put events in queue."""
        import time
        agent_logger.info(f"Task started for agent: {agent.agent_id}")
        try:
            event_count = 0
            async for event in agent.run():
                event_count += 1
                await self.event_queue.put(event)
                if event_count % 100 == 0:
                    agent_logger.debug(f"{agent.agent_id}: {event_count} events emitted")
            agent_logger.info(f"Task completed for agent: {agent.agent_id} | events={event_count}")
        except Exception as e:
            import traceback
            agent_logger.error(f"Task error for agent {agent.agent_id}: {e}")
            agent_logger.error(traceback.format_exc())
            await self.event_queue.put(AgentEvent(
                agent.agent_id, "error", {"error": str(e)}
            ))

    async def _run_sequential_agent(self, agent: BaseAgent) -> tuple[list[dict] | None, str]:
        """Run a sequential agent (DataBinding or UX) and collect results."""
        patches = None
        raw = ""
        async for event in agent.run():
            if event.event_type == "token":
                raw = event.data.get("full", "")
            elif event.event_type == "complete":
                patches = event.data.get("patches", [])
                raw = event.data.get("raw", "")
            elif event.event_type == "error":
                agent_logger.error(f"Agent {agent.agent_id} error: {event.data.get('error')}")
        return patches, raw

    async def _expand_placeholders_recursive(
        self,
        schema: dict,
        placeholders: list[dict],
        prompt: str,
        current_depth: int = 1
    ) -> AsyncGenerator[dict, None]:
        """
        Recursively expand placeholders at the current depth level.
        Returns events and the updated schema.
        """
        if not placeholders or current_depth > self.max_depth:
            return

        depth_label = f"Depth {current_depth}/{self.max_depth}"
        yield {"event": "trace", "data": {
            "type": "agent_start",
            "content": f"{depth_label}: Expanding {len(placeholders)} placeholders..."
        }}
        yield {"event": "phase", "data": {
            "phase": "expansion",
            "status": "started",
            "depth": current_depth,
            "total": len(placeholders)
        }}
        yield {"event": "placeholders", "data": [
            {"id": p["id"], "description": p["description"], "depth": current_depth}
            for p in placeholders
        ]}

        # Get context summary from store
        context_summary = self.store.get_context_summary() if self.store else ""

        # Create and launch expansion agents in parallel
        tasks = []
        for i, placeholder in enumerate(placeholders):
            agent_model = EXPANSION_MODELS[i % len(EXPANSION_MODELS)]
            agent_logger.info(f"Creating ExpansionAgent for {placeholder['id']} at depth {current_depth}")

            agent = ExpansionAgent(
                agent_id=f"expansion_d{current_depth}_{i}_{placeholder['id']}",
                client=self.client,
                model=agent_model,
                placeholder=placeholder,
                overall_purpose=prompt,
                current_depth=current_depth,
                max_depth=self.max_depth,
                context_summary=context_summary
            )

            task = asyncio.create_task(self._run_agent_and_queue_events(agent))
            tasks.append(task)

            yield {"event": "expansion_start", "data": {
                "placeholder_id": placeholder["id"],
                "agent_id": agent.agent_id,
                "model": agent_model,
                "depth": current_depth
            }}

        # Collect events from all agents
        completed_count = 0
        total_placeholders = len(placeholders)
        partial_expansions = {}
        new_placeholders = []  # Collect new placeholders for next depth

        while completed_count < total_placeholders:
            try:
                event = await asyncio.wait_for(self.event_queue.get(), timeout=0.1)

                if event.event_type == "token":
                    placeholder_id = event.data["placeholder_id"]
                    parsed = event.data.get("parsed")
                    token = event.data.get("token", "")
                    full = event.data.get("full", "")

                    yield {"event": "json_chunk", "data": {
                        "chunk": token,
                        "full": full,
                        "closed": json.dumps(parsed, ensure_ascii=False) if parsed else None,
                        "placeholder_id": placeholder_id,
                        "depth": current_depth
                    }}

                    if parsed:
                        partial_expansions[placeholder_id] = parsed

                elif event.event_type == "complete":
                    completed_count += 1
                    placeholder_id = event.data["placeholder_id"]
                    expanded = event.data["expanded"]
                    path = event.data["path"]

                    self.expanded_results[placeholder_id] = expanded
                    self.total_expansions += 1

                    # Save fragment to store
                    if self.store and expanded:
                        self.store.save_fragment(placeholder_id, expanded, {
                            "depth": current_depth,
                            "path": path,
                            "parent": placeholder.get("parent_id")
                        })

                    # Check for new placeholders in the expanded content
                    if expanded and current_depth < self.max_depth:
                        sub_placeholders = find_placeholders(expanded)
                        if sub_placeholders:
                            # Update paths to be relative to current placeholder
                            for sp in sub_placeholders:
                                sp["path"] = f"{path}{sp['path']}"
                                sp["parent_id"] = placeholder_id
                            new_placeholders.extend(sub_placeholders)
                            agent_logger.info(f"Found {len(sub_placeholders)} new placeholders in {placeholder_id}")

                    yield {"event": "expansion_complete", "data": {
                        "placeholder_id": placeholder_id,
                        "success": expanded is not None,
                        "completed": completed_count,
                        "total": total_placeholders,
                        "depth": current_depth,
                        "has_sub_placeholders": len(find_placeholders(expanded) if expanded else []) > 0
                    }}

                    yield {"event": "trace", "data": {
                        "type": "agent_complete",
                        "content": f"{depth_label}: {placeholder_id} ({completed_count}/{total_placeholders})"
                    }}

                elif event.event_type == "error":
                    completed_count += 1
                    yield {"event": "expansion_error", "data": {**event.data, "depth": current_depth}}

            except asyncio.TimeoutError:
                if all(task.done() for task in tasks):
                    while not self.event_queue.empty():
                        event = await self.event_queue.get()
                        if event.event_type == "complete":
                            completed_count += 1
                        elif event.event_type == "error":
                            completed_count += 1
                    break
                continue

        await asyncio.gather(*tasks, return_exceptions=True)

        yield {"event": "phase", "data": {
            "phase": "expansion",
            "status": "completed",
            "depth": current_depth
        }}

        # Store new placeholders for recursive expansion
        self.current_depth_expansions[current_depth] = {
            "expanded": self.expanded_results.copy(),
            "new_placeholders": new_placeholders
        }

        # Yield info about next depth if there are new placeholders
        if new_placeholders and current_depth < self.max_depth:
            yield {"event": "trace", "data": {
                "type": "thinking",
                "content": f"Found {len(new_placeholders)} sub-placeholders, will expand at depth {current_depth + 1}"
            }}
            yield {"event": "new_placeholders", "data": {
                "count": len(new_placeholders),
                "next_depth": current_depth + 1,
                "placeholders": [{"id": p["id"], "description": p["description"][:100]} for p in new_placeholders]
            }}

    async def run(self, prompt: str) -> AsyncGenerator[dict, None]:
        """
        Run hierarchical generation with recursive expansion.

        Flow:
        1. SkeletonAgent generates structure (streaming)
        2. Find all placeholders
        3. Recursively expand placeholders (depth-first, parallel at each level)
        4. DataBindingAgent adds data to components
        5. UXEnhancementAgent improves interactions
        6. Save and return final schema
        """
        import time
        start_time = time.time()
        session_id = str(uuid.uuid4())[:8]

        # Initialize schema store
        self.store = SchemaStore(session_id)
        self.store.save_session_meta({
            "prompt": prompt,
            "max_depth": self.max_depth,
            "created_at": datetime.now().isoformat(),
            "status": "started"
        })

        logger.info(f"Starting recursive generation | session={session_id} | max_depth={self.max_depth}")
        llm_logger.info(f"SESSION_START | session={session_id} | mode=recursive | max_depth={self.max_depth}")
        llm_logger.info(f"SESSION_PROMPT | session={session_id} | prompt={prompt}")

        yield {"event": "start", "data": {
            "session_id": session_id,
            "mode": "recursive",
            "max_depth": self.max_depth
        }}
        yield {"event": "trace", "data": {"type": "thinking", "content": f"Analyzing request: {prompt}"}}

        # =====================================================================
        # Phase 1: Run SkeletonAgent
        # =====================================================================
        yield {"event": "trace", "data": {"type": "agent_start", "content": "Phase 1: Starting SkeletonAgent..."}}
        yield {"event": "phase", "data": {"phase": "skeleton", "status": "started"}}

        skeleton_agent = SkeletonAgent(
            agent_id="skeleton",
            client=self.client,
            model=self.model,
            prompt=prompt
        )

        skeleton = None
        skeleton_raw = ""

        try:
            async for event in skeleton_agent.run():
                if event.event_type == "token":
                    skeleton_raw = event.data["full"]
                    parsed = event.data.get("parsed")
                    if parsed:
                        preview_schema = convert_placeholders_to_skeleton(parsed)
                        yield {"event": "partial_schema", "data": preview_schema}
                    yield {"event": "json_chunk", "data": {
                        "chunk": event.data["token"],
                        "full": skeleton_raw,
                        "closed": None
                    }}
                elif event.event_type == "complete":
                    skeleton = event.data["skeleton"]
                    skeleton_raw = event.data["raw"]
                elif event.event_type == "error":
                    yield {"event": "error", "data": {"error": event.data["error"]}}
                    return

            if not skeleton:
                yield {"event": "error", "data": {"error": "SkeletonAgent failed to generate skeleton"}}
                return

            # Save skeleton to store (both original and with refs)
            self.store.save_skeleton(skeleton)
            skeleton_refs = self.store.save_skeleton_with_refs(skeleton)

            yield {"event": "phase", "data": {"phase": "skeleton", "status": "completed"}}
            yield {"event": "skeleton", "data": skeleton}
            yield {"event": "skeleton_refs", "data": skeleton_refs}
            yield {"event": "trace", "data": {"type": "agent_complete", "content": "SkeletonAgent completed"}}

            # Find initial placeholders
            placeholders = find_placeholders(skeleton)
            agent_logger.info(f"Found {len(placeholders)} initial placeholders")

            current_schema = skeleton

            # =====================================================================
            # Phase 2: Recursive Expansion
            # =====================================================================
            current_depth = 1
            current_placeholders = placeholders

            while current_placeholders and current_depth <= self.max_depth:
                yield {"event": "trace", "data": {
                    "type": "thinking",
                    "content": f"Starting expansion at depth {current_depth}/{self.max_depth} ({len(current_placeholders)} placeholders)"
                }}

                # Expand all placeholders at current depth
                async for event in self._expand_placeholders_recursive(
                    current_schema,
                    current_placeholders,
                    prompt,
                    current_depth
                ):
                    yield event

                # Apply all expansions to current schema
                for placeholder in current_placeholders:
                    expanded = self.expanded_results.get(placeholder["id"])
                    if expanded:
                        agent_logger.info(f"Replacing placeholder {placeholder['id']} at path {placeholder['path']}")
                        current_schema = replace_placeholder(current_schema, placeholder["path"], expanded)
                    else:
                        agent_logger.warning(f"No expansion found for placeholder {placeholder['id']}")

                # Update preview
                preview_schema = convert_placeholders_to_skeleton(current_schema)
                yield {"event": "partial_schema", "data": preview_schema}

                # Get new placeholders for next depth
                depth_data = self.current_depth_expansions.get(current_depth, {})
                current_placeholders = depth_data.get("new_placeholders", [])
                current_depth += 1

            yield {"event": "trace", "data": {
                "type": "agent_complete",
                "content": f"Recursive expansion complete: {self.total_expansions} total expansions across {current_depth - 1} depth levels"
            }}

            # =====================================================================
            # Phase 3: DataBindingAgent - Add data to components (optional)
            # =====================================================================
            if self.enable_data_binding:
                yield {"event": "trace", "data": {"type": "agent_start", "content": "Phase 3: Starting DataBindingAgent..."}}
                yield {"event": "phase", "data": {"phase": "data_binding", "status": "started"}}

                data_binding_agent = DataBindingAgent(
                    agent_id="data_binding",
                    client=self.client,
                    model=self.model,
                    schema=current_schema,
                    overall_purpose=prompt
                )

                data_patches, data_raw = await self._run_sequential_agent(data_binding_agent)

                if data_patches:
                    agent_logger.info(f"DataBindingAgent generated {len(data_patches)} patches")
                    try:
                        current_schema = apply_json_patch(current_schema, data_patches)
                        yield {"event": "partial_schema", "data": current_schema}
                        yield {"event": "trace", "data": {
                            "type": "agent_complete",
                            "content": f"DataBindingAgent applied {len(data_patches)} data bindings"
                        }}
                    except Exception as e:
                        agent_logger.warning(f"DataBindingAgent patch error: {e}")
                        yield {"event": "trace", "data": {
                            "type": "warning",
                            "content": f"Some data bindings failed: {e}"
                        }}
                else:
                    yield {"event": "trace", "data": {
                        "type": "agent_complete",
                        "content": "DataBindingAgent: No data bindings needed"
                    }}

                yield {"event": "phase", "data": {"phase": "data_binding", "status": "completed"}}
            else:
                yield {"event": "trace", "data": {"type": "skipped", "content": "Phase 3: DataBindingAgent skipped (disabled)"}}
                yield {"event": "phase", "data": {"phase": "data_binding", "status": "skipped"}}

            # =====================================================================
            # Phase 4: UXEnhancementAgent - Improve interactions (optional)
            # =====================================================================
            if self.enable_ux_enhancement:
                yield {"event": "trace", "data": {"type": "agent_start", "content": "Phase 4: Starting UXEnhancementAgent..."}}
                yield {"event": "phase", "data": {"phase": "ux_enhancement", "status": "started"}}

                ux_agent = UXEnhancementAgent(
                    agent_id="ux_enhancement",
                    client=self.client,
                    model=self.model,
                    schema=current_schema,
                    overall_purpose=prompt
                )

                ux_patches, ux_raw = await self._run_sequential_agent(ux_agent)

                if ux_patches:
                    agent_logger.info(f"UXEnhancementAgent generated {len(ux_patches)} patches")
                    try:
                        current_schema = apply_json_patch(current_schema, ux_patches)
                        yield {"event": "partial_schema", "data": current_schema}
                        yield {"event": "trace", "data": {
                            "type": "agent_complete",
                            "content": f"UXEnhancementAgent applied {len(ux_patches)} UX improvements"
                        }}
                    except Exception as e:
                        agent_logger.warning(f"UXEnhancementAgent patch error: {e}")
                        yield {"event": "trace", "data": {
                            "type": "warning",
                            "content": f"Some UX enhancements failed: {e}"
                        }}
                else:
                    yield {"event": "trace", "data": {
                        "type": "agent_complete",
                        "content": "UXEnhancementAgent: No UX enhancements needed"
                    }}

                yield {"event": "phase", "data": {"phase": "ux_enhancement", "status": "completed"}}
            else:
                yield {"event": "trace", "data": {"type": "skipped", "content": "Phase 4: UXEnhancementAgent skipped (disabled)"}}
                yield {"event": "phase", "data": {"phase": "ux_enhancement", "status": "skipped"}}

            # =====================================================================
            # Final: Save and return completed schema
            # =====================================================================
            elapsed = time.time() - start_time

            # Save final merged schema
            self.store.save_merged(current_schema)
            self.store.save_session_meta({
                "prompt": prompt,
                "max_depth": self.max_depth,
                "created_at": self.store.load_session_meta().get("created_at"),
                "completed_at": datetime.now().isoformat(),
                "status": "completed",
                "total_expansions": self.total_expansions,
                "elapsed_seconds": elapsed
            })

            yield {"event": "trace", "data": {"type": "merging", "content": "Assembling final UI..."}}
            # Check if schema still has placeholders
            remaining_placeholders = find_placeholders(current_schema)
            if remaining_placeholders:
                agent_logger.warning(f"Final schema still has {len(remaining_placeholders)} unresolved placeholders: {[p['id'] for p in remaining_placeholders]}")
                # Convert remaining placeholders to error indicators so they render visibly
                current_schema = convert_placeholders_to_error(current_schema)
            else:
                agent_logger.info("Final schema has no remaining placeholders - all expanded successfully")
            yield {"event": "schema", "data": current_schema}
            yield {"event": "trace", "data": {
                "type": "complete",
                "content": f"Recursive generation complete ({self.total_expansions} expansions, {elapsed:.1f}s)"
            }}

            logger.info(f"Recursive generation COMPLETED | session={session_id} | expansions={self.total_expansions} | time={elapsed:.1f}s")
            llm_logger.info(f"SESSION_COMPLETE | session={session_id} | expansions={self.total_expansions} | time={elapsed:.1f}s")

            yield {"event": "complete", "data": {
                "schema": current_schema,
                "session_id": session_id,
                "completed": True,
                "total_expansions": self.total_expansions,
                "max_depth_reached": current_depth - 1,
                "phases_completed": ["skeleton", "expansion", "data_binding", "ux_enhancement"],
                "store_path": str(self.store.session_dir)
            }}

        except Exception as e:
            import traceback
            logger.error(f"Recursive generation ERROR | session={session_id} | error={e}")
            logger.error(traceback.format_exc())
            llm_logger.error(f"SESSION_ERROR | session={session_id} | error={e}")

            if self.store:
                self.store.save_session_meta({
                    "prompt": prompt,
                    "status": "error",
                    "error": str(e)
                })

            yield {"event": "trace", "data": {"type": "error", "content": str(e)}}
            yield {"event": "error", "data": {"error": str(e)}}

STREAMING_SYSTEM_PROMPT = """You are a UI schema generator. Generate JSON schemas that render as React UI components.

## CRITICAL: Close-First-Then-Fill Principle (先闭合再填充)
You MUST follow this generation order:
1. First output the COMPLETE CLOSED structure (all brackets matched)
2. Fill in content values AFTER the structure is complete
3. NEVER leave brackets unclosed - always complete each object/array before starting the next

## Generation Order (IMPORTANT!)
When generating JSON, follow this exact order for EACH component:
1. Open object {{
2. Write "type": "ComponentName"
3. Write "props": {{...}} (complete and close it immediately)
4. Write "children": and either:
   - A string value "text content"
   - An array [...] with each child FULLY CLOSED before the next
5. Close object }}

Example of correct generation order:
{{"type":"Card","props":{{}},children":[{{"type":"Text","children":"Hello"}}]}}

## Output Format
Return ONLY a valid JSON object (no markdown, no explanation):
{{
  "type": "ComponentName",
  "props": {{ }},
  "children": [ ]
}}

## Available Components

### Layout
- Col: Vertical container. Props: gap ("sm"|"md"|"lg"), align, justify. Children: array
- Row: Horizontal container. Props: gap, align, justify, wrap. Children: array
- Card: Card container. Children: [CardHeader, CardContent, CardFooter]
- CardHeader, CardContent, CardFooter: Card sections
- CardTitle, CardDescription: Card text elements

### Navigation (IMPORTANT: Use Tabs for any sidebar/navigation pattern!)
- Tabs: Tab container with switchable content. Props: defaultValue (required, string), orientation ("horizontal"|"vertical")
- TabsList: Container for tab triggers. Props: none
- TabsTrigger: Clickable tab button. Props: value (required, must match TabsContent value)
- TabsContent: Content panel shown when tab is active. Props: value (required, must match TabsTrigger value)

### Basic
- Text: Text display. Props: as ("p"|"h1"|"h2"|"h3"|"span"), variant ("default"|"muted"|"destructive"), size ("sm"|"base"|"lg"|"xl"|"2xl"|"3xl")
- Button: Button. Props: variant ("default"|"outline"|"secondary"|"destructive"|"ghost"), size ("sm"|"default"|"lg")
- Badge: Status badge. Props: variant ("default"|"secondary"|"destructive"|"success"|"warning"|"outline")
- Avatar: User avatar. Props: src, alt, fallback, size ("sm"|"default"|"lg"), status ("online"|"offline"|"away")
- Separator: Divider line. Props: orientation ("horizontal"|"vertical")

### Form
- Input: Text input. Props: placeholder, type ("text"|"email"|"password"|"number")
- Textarea: Multi-line input. Props: placeholder, rows
- Select, SelectTrigger, SelectValue, SelectContent, SelectItem: Dropdown select
- Checkbox: Checkbox. Props: label, description
- Switch: Toggle switch. Props: label, description
- Label: Form label

### Feedback
- Alert: Alert message. Props: variant ("default"|"info"|"success"|"warning"|"destructive"), title
- Progress: Progress bar. Props: value (0-100), variant, showValue, label
- Skeleton: Loading placeholder. Props: variant ("default"|"circular"|"text"), width, height

### Data
- Table, TableHeader, TableBody, TableRow, TableHead, TableCell: Data table

## Navigation Pattern (MUST USE for sidebar/navigation)
For any UI with sidebar navigation or switchable sections, ALWAYS use Tabs with vertical orientation:
{{
  "type": "Tabs",
  "props": {{"defaultValue": "tab1", "orientation": "vertical"}},
  "children": [
    {{"type": "TabsList", "children": [
      {{"type": "TabsTrigger", "props": {{"value": "tab1"}}, "children": "Section 1"}},
      {{"type": "TabsTrigger", "props": {{"value": "tab2"}}, "children": "Section 2"}}
    ]}},
    {{"type": "TabsContent", "props": {{"value": "tab1"}}, "children": [...]}},
    {{"type": "TabsContent", "props": {{"value": "tab2"}}, "children": [...]}}
  ]
}}

## Example
User: "Create a login form"
{{
  "type": "Card",
  "children": [
    {{"type": "CardHeader", "children": [
      {{"type": "CardTitle", "children": "Login"}},
      {{"type": "CardDescription", "children": "Enter your credentials"}}
    ]}},
    {{"type": "CardContent", "children": [
      {{"type": "Col", "props": {{"gap": "md"}}, "children": [
        {{"type": "Col", "props": {{"gap": "sm"}}, "children": [
          {{"type": "Label", "children": "Email"}},
          {{"type": "Input", "props": {{"type": "email", "placeholder": "you@example.com"}}}}
        ]}},
        {{"type": "Col", "props": {{"gap": "sm"}}, "children": [
          {{"type": "Label", "children": "Password"}},
          {{"type": "Input", "props": {{"type": "password", "placeholder": "••••••••"}}}}
        ]}}
      ]}}
    ]}},
    {{"type": "CardFooter", "children": [
      {{"type": "Button", "children": "Sign In"}}
    ]}}
  ]
}}

Now generate the UI schema for the user's request. Output ONLY the JSON, nothing else."""


# =============================================================================
# HIERARCHICAL GENERATION PROMPTS
# =============================================================================

SKELETON_SYSTEM_PROMPT = """You are a UI architect. Generate a HIGH-LEVEL skeleton structure for complex UIs.

## Output Format
Return a JSON object with $placeholder markers for complex sections:
{
  "type": "ComponentName",
  "props": { ... },
  "children": [
    {
      "$placeholder": "unique_id",
      "$description": "Detailed description of what this section should contain"
    }
  ]
}

## Rules
1. Use $placeholder for sections that need detailed content (3+ components)
2. Simple elements (single Text, Button, etc.) can be inline
3. Each $placeholder must have a unique id and clear $description
4. $description should be detailed enough for another AI to implement
5. Keep the skeleton shallow (max 2-3 levels before placeholder)
6. IMPORTANT: For sidebar navigation, ALWAYS use Tabs with orientation="vertical"

## Available Layout Components
- Row: Horizontal container. Props: gap ("sm"|"md"|"lg")
- Col: Vertical container. Props: gap ("sm"|"md"|"lg")
- Card, CardHeader, CardContent, CardFooter: Card sections
- Tabs: Tab container with switchable content. Props: defaultValue (required), orientation ("horizontal"|"vertical")
- TabsList: Container for tab triggers
- TabsTrigger: Clickable tab button. Props: value (required)
- TabsContent: Content panel. Props: value (required)

## CRITICAL: Sidebar Navigation Pattern
For ANY UI with sidebar navigation, you MUST use Tabs with vertical orientation:
{
  "type": "Tabs",
  "props": {"defaultValue": "first_tab", "orientation": "vertical"},
  "children": [
    {"type": "TabsList", "children": [
      {"type": "TabsTrigger", "props": {"value": "tab1"}, "children": "Tab 1"},
      {"type": "TabsTrigger", "props": {"value": "tab2"}, "children": "Tab 2"}
    ]},
    {"type": "TabsContent", "props": {"value": "tab1"}, "children": [{"$placeholder": "...", "$description": "..."}]},
    {"type": "TabsContent", "props": {"value": "tab2"}, "children": [{"$placeholder": "...", "$description": "..."}]}
  ]
}

## Example
User: "Create a settings page with sidebar navigation and main content area"
{
  "type": "Tabs",
  "props": {"defaultValue": "profile", "orientation": "vertical"},
  "children": [
    {"type": "TabsList", "children": [
      {"type": "TabsTrigger", "props": {"value": "profile"}, "children": "Profile"},
      {"type": "TabsTrigger", "props": {"value": "security"}, "children": "Security"},
      {"type": "TabsTrigger", "props": {"value": "notifications"}, "children": "Notifications"},
      {"type": "TabsTrigger", "props": {"value": "billing"}, "children": "Billing"}
    ]},
    {"type": "TabsContent", "props": {"value": "profile"}, "children": [
      {
        "$placeholder": "profile_content",
        "$description": "Profile settings with name input, email input, bio textarea, and save button."
      }
    ]},
    {"type": "TabsContent", "props": {"value": "security"}, "children": [
      {
        "$placeholder": "security_content",
        "$description": "Security settings with password change form and two-factor authentication toggle."
      }
    ]},
    {"type": "TabsContent", "props": {"value": "notifications"}, "children": [
      {
        "$placeholder": "notifications_content",
        "$description": "Notification preferences with email and push notification toggles."
      }
    ]},
    {"type": "TabsContent", "props": {"value": "billing"}, "children": [
      {
        "$placeholder": "billing_content",
        "$description": "Billing information with payment method and subscription details."
      }
    ]}
  ]
}

Now generate the skeleton for the user's request. Output ONLY the JSON."""


EXPAND_PLACEHOLDER_PROMPT = """You are a UI component generator. Expand a placeholder into detailed UI components.

## Context
You are expanding a section of a larger UI. Here is the context:
- Overall UI purpose: {overall_purpose}
- Section to expand: {placeholder_id}
- Section description: {placeholder_description}

## Output Format
Return ONLY a valid JSON object (the expanded UI section):
{
  "type": "ComponentName",
  "props": { ... },
  "children": [ ... ]
}

## Available Components

### Layout
- Col: Vertical container. Props: gap ("sm"|"md"|"lg"), align, justify
- Row: Horizontal container. Props: gap, align, justify, wrap
- Card, CardHeader, CardContent, CardFooter, CardTitle, CardDescription

### Navigation
- Tabs: Tab container. Props: defaultValue (required), orientation ("horizontal"|"vertical")
- TabsList: Container for tab triggers
- TabsTrigger: Tab button. Props: value (required)
- TabsContent: Tab panel. Props: value (required)

### Basic
- Text: Props: as ("p"|"h1"|"h2"|"h3"), variant ("default"|"muted"), size ("sm"|"base"|"lg"|"xl"|"2xl")
- Button: Props: variant ("default"|"outline"|"secondary"|"destructive"|"ghost"), size ("sm"|"default"|"lg")
- Badge: Props: variant ("default"|"secondary"|"destructive"|"success"|"warning")
- Avatar: Props: src, alt, fallback, size ("sm"|"default"|"lg")
- Separator: Props: orientation ("horizontal"|"vertical")

### Form
- Input: Props: placeholder, type ("text"|"email"|"password"|"number")
- Textarea: Props: placeholder, rows
- Select, SelectTrigger, SelectValue, SelectContent, SelectItem
- Checkbox: Props: label, description
- Switch: Props: label, description
- Label: Form label
- Slider: Props: min, max, step

### Feedback
- Alert: Props: variant ("default"|"info"|"success"|"warning"|"destructive"), title
- Progress: Props: value (0-100), showValue, label

### Data
- Table, TableHeader, TableBody, TableRow, TableHead, TableCell
- Statistic: Props: title, value, prefix, suffix, trend ("up"|"down"|"neutral"), trendValue, description

### Charts
- LineChart, BarChart, PieChart, AreaChart: Props: data (array), xKey, yKey, title

Now expand the placeholder into detailed components. Output ONLY the JSON."""


def get_component_docs_for_category(categories: list[str]) -> str:
    """Read component documentation for specified categories."""
    docs = []
    for category in categories:
        readme_path = COMPONENTS_DIR / category / "README.md"
        if readme_path.exists():
            docs.append(f"## {category.title()} Components\n{readme_path.read_text()}")
    return "\n\n".join(docs)


MODIFY_STREAMING_PROMPT = """You are a UI schema modifier. Output JSON Patch operations to modify the schema.

## Output Format
Output ONLY a JSON array of patch operations:
[
  {{"op": "replace", "path": "/path/to/value", "value": "new value"}}
]

## Operations
- "replace": Change existing value
- "add": Add new property or array item
- "remove": Delete property or array item

## Path Format (JSON Pointer)
- Start with /
- Use property names: /props/variant
- Use array indices (0-based): /children/0
- Combine: /children/0/props/placeholder

## Current Schema Structure
{schema}

## Schema Path Reference
{path_reference}

## Modification Request
{prompt}

## Examples
Change text: [{{"op":"replace","path":"/children/0/children","value":"New Text"}}]
Change prop: [{{"op":"replace","path":"/children/1/props/variant","value":"destructive"}}]
Add prop: [{{"op":"add","path":"/children/0/props/disabled","value":true}}]

Output ONLY the JSON array:"""


def generate_path_reference(schema: dict, prefix: str = "") -> list[str]:
    """Generate a list of all paths in the schema for LLM reference."""
    paths = []

    if isinstance(schema, dict):
        for key, value in schema.items():
            if key == '_id':
                continue
            current_path = f"{prefix}/{key}"
            if key == 'type':
                paths.append(f"{current_path} = \"{value}\"")
            elif key == 'children':
                if isinstance(value, str):
                    paths.append(f"{current_path} = \"{value}\"")
                elif isinstance(value, list):
                    for i, child in enumerate(value):
                        if isinstance(child, str):
                            paths.append(f"{current_path}/{i} = \"{child}\"")
                        else:
                            child_type = child.get('type', 'unknown')
                            paths.append(f"{current_path}/{i} (type: {child_type})")
                            paths.extend(generate_path_reference(child, f"{current_path}/{i}"))
            elif key == 'props' and isinstance(value, dict):
                for prop_key, prop_value in value.items():
                    paths.append(f"{current_path}/{prop_key} = {json.dumps(prop_value)}")
            elif isinstance(value, dict):
                paths.extend(generate_path_reference(value, current_path))

    return paths


def read_component_docs(category: str) -> str:
    """Read component documentation from README files."""
    readme_path = COMPONENTS_DIR / category / "README.md"
    if readme_path.exists():
        return readme_path.read_text()
    return f"No documentation found for category: {category}"


def parse_partial_json(json_str: str) -> dict | None:
    """
    Try to parse partial JSON by auto-closing brackets and strings.
    Optimized for streaming: tries to produce valid JSON on every token.
    """
    if not json_str.strip():
        return None

    # Clean up the string
    s = json_str.strip()

    # Remove markdown code blocks if present
    if s.startswith("```"):
        lines = s.split("\n")
        s = "\n".join(lines[1:])
        if "```" in s:
            s = s[:s.rfind("```")]

    # Find the start of JSON
    start = s.find("{")
    if start == -1:
        return None
    s = s[start:]

    # Try parsing as-is first
    try:
        return json.loads(s)
    except json.JSONDecodeError:
        pass

    # Smart closing: track state and build closable string
    in_string = False
    escape_next = False
    stack = []  # Track open braces/brackets with their positions
    last_complete_pos = 0  # Position after last complete value
    colon_pos = -1  # Position of last colon outside string
    comma_pos = -1  # Position of last comma outside string

    for i, char in enumerate(s):
        if escape_next:
            escape_next = False
            continue

        if char == '\\' and in_string:
            escape_next = True
            continue

        if char == '"':
            in_string = not in_string
            if not in_string:
                # Just closed a string - this is a complete value
                last_complete_pos = i + 1
            continue

        if in_string:
            continue

        # Not in string
        if char == '{':
            stack.append(('}', i))
        elif char == '[':
            stack.append((']', i))
        elif char == '}':
            if stack and stack[-1][0] == '}':
                stack.pop()
                last_complete_pos = i + 1
        elif char == ']':
            if stack and stack[-1][0] == ']':
                stack.pop()
                last_complete_pos = i + 1
        elif char == ':':
            colon_pos = i
        elif char == ',':
            comma_pos = i
            last_complete_pos = i + 1

    # Strategy 1: Close string and brackets directly
    closed = s
    if in_string:
        closed += '"'

    # Remove trailing colon or comma
    closed = re.sub(r'[,:]\s*$', '', closed)

    # Close all open brackets
    for closer, _ in reversed(stack):
        closed += closer

    try:
        return json.loads(closed)
    except json.JSONDecodeError:
        pass

    # Strategy 2: Truncate to last complete position and close
    if last_complete_pos > 0:
        truncated = s[:last_complete_pos]
        # Remove trailing comma
        truncated = re.sub(r',\s*$', '', truncated)
        # Close brackets
        open_braces = truncated.count('{') - truncated.count('}')
        open_brackets = truncated.count('[') - truncated.count(']')
        truncated += ']' * max(0, open_brackets)
        truncated += '}' * max(0, open_braces)
        try:
            return json.loads(truncated)
        except json.JSONDecodeError:
            pass

    # Strategy 3: Find last comma and truncate there
    if comma_pos > 0:
        truncated = s[:comma_pos]
        open_braces = truncated.count('{') - truncated.count('}')
        open_brackets = truncated.count('[') - truncated.count(']')
        truncated += ']' * max(0, open_brackets)
        truncated += '}' * max(0, open_braces)
        try:
            return json.loads(truncated)
        except json.JSONDecodeError:
            pass

    return None


def parse_json_pointer(pointer: str) -> list[str | int]:
    """Parse JSON Pointer (RFC 6901) into path segments."""
    if not pointer:
        return []
    if not pointer.startswith('/'):
        return []
    parts = pointer[1:].split('/')
    result = []
    for part in parts:
        # Unescape ~1 -> / and ~0 -> ~
        part = part.replace('~1', '/').replace('~0', '~')
        # Try to convert to int for array indices
        try:
            result.append(int(part))
        except ValueError:
            result.append(part)
    return result


def get_by_pointer(obj: dict | list, pointer: str):
    """Get value at JSON Pointer path."""
    parts = parse_json_pointer(pointer)
    current = obj
    for part in parts:
        if isinstance(current, dict):
            current = current[part]
        elif isinstance(current, list):
            current = current[part]
        else:
            raise KeyError(f"Cannot traverse {type(current)} with {part}")
    return current


def set_by_pointer(obj: dict | list, pointer: str, value) -> dict | list:
    """Set value at JSON Pointer path. Returns modified copy."""
    import copy
    obj = copy.deepcopy(obj)
    parts = parse_json_pointer(pointer)
    if not parts:
        return value

    current = obj
    for part in parts[:-1]:
        if isinstance(current, dict):
            current = current[part]
        elif isinstance(current, list):
            current = current[part]

    last_part = parts[-1]
    if isinstance(current, dict):
        current[last_part] = value
    elif isinstance(current, list):
        if last_part == '-':
            current.append(value)
        else:
            current[last_part] = value

    return obj


def remove_by_pointer(obj: dict | list, pointer: str) -> dict | list:
    """Remove value at JSON Pointer path. Returns modified copy."""
    import copy
    obj = copy.deepcopy(obj)
    parts = parse_json_pointer(pointer)
    if not parts:
        return obj

    current = obj
    for part in parts[:-1]:
        if isinstance(current, dict):
            current = current[part]
        elif isinstance(current, list):
            current = current[part]

    last_part = parts[-1]
    if isinstance(current, dict):
        del current[last_part]
    elif isinstance(current, list):
        del current[last_part]

    return obj


def apply_json_patch(obj: dict, patches: list[dict]) -> dict:
    """Apply JSON Patch operations to object."""
    result = obj
    for patch in patches:
        op = patch.get('op')
        path = patch.get('path', '')
        value = patch.get('value')

        try:
            if op == 'replace':
                result = set_by_pointer(result, path, value)
            elif op == 'add':
                result = set_by_pointer(result, path, value)
            elif op == 'remove':
                result = remove_by_pointer(result, path)
        except (KeyError, IndexError, TypeError) as e:
            # Skip invalid patches
            logger.warning(f"Patch error: {e} for {patch}")
            continue

    return result


def parse_partial_patches(json_str: str) -> list[dict] | None:
    """Parse partial JSON array of patches."""
    if not json_str.strip():
        return None

    s = json_str.strip()

    # Find array start
    start = s.find('[')
    if start == -1:
        return None
    s = s[start:]

    # Try parsing as-is
    try:
        return json.loads(s)
    except json.JSONDecodeError:
        pass

    # Try to close the array
    # Find complete objects within the array
    patches = []
    depth = 0
    in_string = False
    escape_next = False
    obj_start = -1

    for i, char in enumerate(s):
        if escape_next:
            escape_next = False
            continue
        if char == '\\' and in_string:
            escape_next = True
            continue
        if char == '"':
            in_string = not in_string
            continue
        if in_string:
            continue

        if char == '{':
            if depth == 1:  # Inside array, starting new object
                obj_start = i
            depth += 1
        elif char == '}':
            depth -= 1
            if depth == 1 and obj_start >= 0:  # Completed an object inside array
                try:
                    obj_str = s[obj_start:i+1]
                    patches.append(json.loads(obj_str))
                except json.JSONDecodeError:
                    pass
                obj_start = -1
        elif char == '[':
            if depth == 0:
                depth = 1
            else:
                depth += 1
        elif char == ']':
            depth -= 1

    return patches if patches else None


# =============================================================================
# HIERARCHICAL GENERATION FUNCTIONS
# =============================================================================

def find_placeholders(schema: dict, path: str = "") -> list[dict]:
    """Find all $placeholder markers in the schema with their paths."""
    placeholders = []

    if isinstance(schema, dict):
        if "$placeholder" in schema:
            placeholders.append({
                "id": schema["$placeholder"],
                "description": schema.get("$description", ""),
                "path": path
            })
        else:
            for key, value in schema.items():
                if key == "children" and isinstance(value, list):
                    for i, child in enumerate(value):
                        placeholders.extend(find_placeholders(child, f"{path}/children/{i}"))
                elif isinstance(value, dict):
                    placeholders.extend(find_placeholders(value, f"{path}/{key}"))

    return placeholders


def replace_placeholder(schema: dict, path: str, replacement: dict) -> dict:
    """Replace a placeholder at the given path with the replacement content."""
    import copy
    schema = copy.deepcopy(schema)

    if not path:
        return replacement

    parts = [p for p in path.split("/") if p]
    current = schema

    for i, part in enumerate(parts[:-1]):
        if part == "children":
            continue
        elif part.isdigit():
            current = current["children"][int(part)]
        else:
            current = current[part]

    # Handle the last part
    last_part = parts[-1]
    if last_part.isdigit():
        current["children"][int(last_part)] = replacement
    else:
        current[last_part] = replacement

    return schema


def generate_skeleton(prompt: str) -> tuple[dict | None, str]:
    """Generate the skeleton structure with placeholders."""
    messages = [
        {"role": "system", "content": SKELETON_SYSTEM_PROMPT},
        {"role": "user", "content": prompt}
    ]

    response = client.chat.completions.create(
        model=MODEL,
        messages=messages,
        temperature=0.7,
    )

    content = response.choices[0].message.content
    skeleton = parse_partial_json(content)
    return skeleton, content


def generate_skeleton_streaming(prompt: str):
    """Generate skeleton with streaming output."""
    messages = [
        {"role": "system", "content": SKELETON_SYSTEM_PROMPT},
        {"role": "user", "content": prompt}
    ]

    stream = client.chat.completions.create(
        model=MODEL,
        messages=messages,
        temperature=0.7,
        stream=True,
    )

    full_content = ""
    for chunk in stream:
        if chunk.choices and chunk.choices[0].delta.content:
            token = chunk.choices[0].delta.content
            full_content += token
            yield {"type": "token", "token": token, "full": full_content}

    skeleton = parse_partial_json(full_content)
    yield {"type": "complete", "skeleton": skeleton, "raw": full_content}


def expand_placeholder(placeholder: dict, overall_purpose: str) -> tuple[str, dict | None, str]:
    """Expand a single placeholder into detailed components."""
    prompt = EXPAND_PLACEHOLDER_PROMPT.format(
        overall_purpose=overall_purpose,
        placeholder_id=placeholder["id"],
        placeholder_description=placeholder["description"]
    )

    messages = [
        {"role": "user", "content": prompt}
    ]

    response = client.chat.completions.create(
        model=MODEL,
        messages=messages,
        temperature=0.7,
    )

    content = response.choices[0].message.content
    expanded = parse_partial_json(content)
    return placeholder["id"], expanded, content


def expand_placeholder_streaming(placeholder: dict, overall_purpose: str):
    """Expand a placeholder with streaming output."""
    prompt = EXPAND_PLACEHOLDER_PROMPT.format(
        overall_purpose=overall_purpose,
        placeholder_id=placeholder["id"],
        placeholder_description=placeholder["description"]
    )

    messages = [
        {"role": "user", "content": prompt}
    ]

    stream = client.chat.completions.create(
        model=MODEL,
        messages=messages,
        temperature=0.7,
        stream=True,
    )

    full_content = ""
    for chunk in stream:
        if chunk.choices and chunk.choices[0].delta.content:
            token = chunk.choices[0].delta.content
            full_content += token
            yield {"type": "token", "placeholder_id": placeholder["id"], "token": token, "full": full_content}

    expanded = parse_partial_json(full_content)
    yield {"type": "complete", "placeholder_id": placeholder["id"], "expanded": expanded, "raw": full_content}


def convert_placeholders_to_skeleton(schema: dict) -> dict:
    """Convert $placeholder markers to Skeleton components for rendering."""
    import copy
    schema = copy.deepcopy(schema)

    def convert_node(node):
        if isinstance(node, dict):
            if "$placeholder" in node:
                # Convert to a visible skeleton placeholder
                return {
                    "type": "Card",
                    "props": {"variant": "outline"},
                    "children": [
                        {
                            "type": "CardContent",
                            "children": [
                                {
                                    "type": "Col",
                                    "props": {"gap": "sm"},
                                    "children": [
                                        {
                                            "type": "Row",
                                            "props": {"gap": "sm", "align": "center"},
                                            "children": [
                                                {"type": "Skeleton", "props": {"variant": "circular", "width": 24, "height": 24}},
                                                {"type": "Text", "props": {"variant": "muted", "size": "sm"}, "children": f"Loading: {node['$placeholder']}..."}
                                            ]
                                        },
                                        {"type": "Skeleton", "props": {"variant": "text", "width": "100%", "height": 16}},
                                        {"type": "Skeleton", "props": {"variant": "text", "width": "80%", "height": 16}},
                                        {"type": "Skeleton", "props": {"variant": "default", "width": "100%", "height": 60}}
                                    ]
                                }
                            ]
                        }
                    ],
                    "_placeholder_id": node["$placeholder"]
                }
            else:
                result = {}
                for key, value in node.items():
                    if key == "children" and isinstance(value, list):
                        result[key] = [convert_node(child) for child in value]
                    elif isinstance(value, dict):
                        result[key] = convert_node(value)
                    else:
                        result[key] = value
                return result
        return node

    return convert_node(schema)


def convert_placeholders_to_error(schema: dict) -> dict:
    """Convert remaining $placeholder markers to error indicator components for final rendering."""
    import copy
    schema = copy.deepcopy(schema)

    def convert_node(node):
        if isinstance(node, dict):
            if "$placeholder" in node:
                # Convert to a visible error indicator
                return {
                    "type": "Card",
                    "props": {"variant": "outline", "className": "border-red-300 bg-red-50"},
                    "children": [
                        {
                            "type": "CardContent",
                            "children": [
                                {
                                    "type": "Col",
                                    "props": {"gap": "sm", "align": "center"},
                                    "children": [
                                        {"type": "Text", "props": {"variant": "muted", "size": "sm", "className": "text-red-600"}, "children": f"⚠️ 展开失败: {node['$placeholder']}"},
                                        {"type": "Text", "props": {"variant": "muted", "size": "xs"}, "children": node.get("$description", "无描述")}
                                    ]
                                }
                            ]
                        }
                    ]
                }
            else:
                result = {}
                for key, value in node.items():
                    if key == "children" and isinstance(value, list):
                        result[key] = [convert_node(child) for child in value]
                    elif isinstance(value, dict):
                        result[key] = convert_node(value)
                    else:
                        result[key] = value
                return result
        return node

    return convert_node(schema)


async def run_hierarchical_generation(
    prompt: str,
    max_depth: int = 3,
    enable_data_binding: bool = True,
    enable_ux_enhancement: bool = True
) -> AsyncGenerator[dict, None]:
    """
    Run hierarchical UI generation using multi-agent architecture with recursive expansion.

    Uses AgentOrchestrator to coordinate:
    1. SkeletonAgent - generates structure
    2. Multiple ExpansionAgents - expand placeholders recursively (up to max_depth)
    3. DataBindingAgent - adds data to components (optional)
    4. UXEnhancementAgent - improves interactions (optional)
    """
    orchestrator = AgentOrchestrator(
        client=client,
        model=MODEL,
        max_depth=max_depth,
        enable_data_binding=enable_data_binding,
        enable_ux_enhancement=enable_ux_enhancement
    )
    async for event in orchestrator.run(prompt):
        yield event


async def run_streaming_generation(prompt: str) -> AsyncGenerator[dict, None]:
    """
    Stream UI generation token by token.
    Yields partial JSON and parsed schema as they become available.
    """
    import time
    start_time = time.time()
    session_id = str(uuid.uuid4())[:8]

    logger.info(f"run_streaming_generation STARTED | session={session_id} | prompt_length={len(prompt)}")

    messages = [
        {"role": "system", "content": STREAMING_SYSTEM_PROMPT},
        {"role": "user", "content": prompt}
    ]

    # Log LLM request
    llm_logger.info(f"LLM_REQUEST | session={session_id} | model={MODEL} | mode=simple_streaming")
    llm_logger.info(f"LLM_USER_PROMPT | session={session_id} | prompt={prompt}")
    llm_logger.debug(f"LLM_SYSTEM_PROMPT | session={session_id} | length={len(STREAMING_SYSTEM_PROMPT)}")

    yield {"event": "start", "data": {"session_id": session_id}}
    yield {"event": "trace", "data": {"type": "thinking", "content": f"Understanding request: {prompt}"}}

    try:
        # Use streaming API
        yield {"event": "trace", "data": {"type": "generating", "content": "Generating UI schema..."}}

        stream = await client.chat.completions.create(
            model=MODEL,
            messages=messages,
            temperature=0.7,
            stream=True,
        )

        full_content = ""
        last_closed_json = ""
        token_count = 0

        async for chunk in stream:
            if chunk.choices and chunk.choices[0].delta.content:
                token = chunk.choices[0].delta.content
                full_content += token
                token_count += 1

                # Log progress every 50 tokens
                if token_count % 50 == 0:
                    elapsed = time.time() - start_time
                    agent_logger.debug(f"simple_streaming: {token_count} tokens ({elapsed:.1f}s)")

                # Try to parse and close partial JSON
                parsed = parse_partial_json(full_content)
                if parsed:
                    closed_json = json.dumps(parsed, indent=2, ensure_ascii=False)
                    # Always send json_chunk when closed JSON changes
                    if closed_json != last_closed_json:
                        last_closed_json = closed_json
                        yield {"event": "json_chunk", "data": {"chunk": token, "full": full_content, "closed": closed_json}}
                        # Also update partial_schema for UI rendering
                        yield {"event": "partial_schema", "data": parsed}
                else:
                    # Send raw content when not parseable
                    yield {"event": "json_chunk", "data": {"chunk": token, "full": full_content, "closed": None}}

        # Final parse
        elapsed = time.time() - start_time
        final_schema = parse_partial_json(full_content)

        # Log complete LLM response
        logger.info(f"run_streaming_generation COMPLETED | session={session_id} | tokens={token_count} | time={elapsed:.1f}s")
        llm_logger.info(f"LLM_RESPONSE | session={session_id} | tokens={token_count} | time={elapsed:.1f}s")
        llm_logger.info(f"LLM_OUTPUT_CONTENT | session={session_id} |\n{full_content}")

        if final_schema:
            yield {"event": "schema", "data": final_schema}
            yield {"event": "trace", "data": {"type": "complete", "content": "UI generation complete"}}

        yield {
            "event": "complete",
            "data": {
                "schema": final_schema,
                "session_id": session_id,
                "completed": final_schema is not None
            }
        }

    except Exception as e:
        yield {"event": "trace", "data": {"type": "error", "content": str(e)}}
        yield {"event": "error", "data": {"error": str(e)}}


async def run_streaming_modification(prompt: str, current_schema: dict) -> AsyncGenerator[dict, None]:
    """
    Stream UI modification using JSON Patch.
    LLM outputs patch operations, we apply them incrementally.
    """
    import time
    start_time = time.time()
    session_id = str(uuid.uuid4())[:8]

    logger.info(f"run_streaming_modification STARTED | session={session_id} | prompt_length={len(prompt)}")

    # Generate path reference for LLM
    path_ref = generate_path_reference(current_schema)
    path_reference_str = "\n".join(path_ref[:50])  # Limit to 50 paths

    schema_json = json.dumps(current_schema, indent=2)
    full_prompt = MODIFY_STREAMING_PROMPT.format(
        schema=schema_json,
        path_reference=path_reference_str,
        prompt=prompt
    )

    messages = [
        {"role": "user", "content": full_prompt}
    ]

    # Log LLM request
    llm_logger.info(f"LLM_REQUEST | session={session_id} | model={MODEL} | mode=modification")
    llm_logger.info(f"LLM_USER_PROMPT | session={session_id} | modification_request={prompt}")
    llm_logger.debug(f"LLM_FULL_PROMPT | session={session_id} | length={len(full_prompt)}")
    llm_logger.debug(f"LLM_CURRENT_SCHEMA | session={session_id} |\n{schema_json[:2000]}...")

    yield {"event": "start", "data": {"session_id": session_id}}
    yield {"event": "trace", "data": {"type": "thinking", "content": f"Understanding modification: {prompt}"}}

    try:
        yield {"event": "trace", "data": {"type": "generating", "content": "Generating patches..."}}

        stream = await client.chat.completions.create(
            model=MODEL,
            messages=messages,
            temperature=0.7,
            stream=True,
        )

        full_content = ""
        last_patches_count = 0
        last_schema_json = ""
        token_count = 0

        async for chunk in stream:
            if chunk.choices and chunk.choices[0].delta.content:
                token = chunk.choices[0].delta.content
                full_content += token
                token_count += 1

                # Log progress every 50 tokens
                if token_count % 50 == 0:
                    elapsed = time.time() - start_time
                    agent_logger.debug(f"modification: {token_count} tokens ({elapsed:.1f}s)")

                # Send raw patch content
                yield {"event": "json_chunk", "data": {"chunk": token, "full": full_content, "closed": full_content}}

                # Try to parse patches and apply them
                patches = parse_partial_patches(full_content)
                if patches and len(patches) > last_patches_count:
                    # New patches available
                    for i in range(last_patches_count, len(patches)):
                        patch = patches[i]
                        yield {"event": "patch", "data": patch}

                    last_patches_count = len(patches)

                    # Apply all patches to get current schema
                    try:
                        modified_schema = apply_json_patch(current_schema, patches)
                        schema_json_str = json.dumps(modified_schema, sort_keys=True)
                        if schema_json_str != last_schema_json:
                            last_schema_json = schema_json_str
                            yield {"event": "partial_schema", "data": modified_schema}
                    except Exception as e:
                        logger.warning(f"Patch apply error: {e}")

        # Final patches
        elapsed = time.time() - start_time
        final_patches = parse_partial_patches(full_content)
        final_schema = current_schema

        # Log complete LLM response
        logger.info(f"run_streaming_modification COMPLETED | session={session_id} | tokens={token_count} | patches={len(final_patches or [])} | time={elapsed:.1f}s")
        llm_logger.info(f"LLM_RESPONSE | session={session_id} | tokens={token_count} | time={elapsed:.1f}s")
        llm_logger.info(f"LLM_OUTPUT_CONTENT | session={session_id} |\n{full_content}")

        if final_patches:
            try:
                final_schema = apply_json_patch(current_schema, final_patches)
            except Exception as e:
                logger.warning(f"Final patch error: {e}")

        yield {"event": "schema", "data": final_schema}
        yield {"event": "trace", "data": {"type": "complete", "content": f"Applied {len(final_patches or [])} patches"}}

        yield {
            "event": "complete",
            "data": {
                "schema": final_schema,
                "session_id": session_id,
                "completed": True
            }
        }

    except Exception as e:
        yield {"event": "trace", "data": {"type": "error", "content": str(e)}}
        yield {"event": "error", "data": {"error": str(e)}}


# Keep the old agent-based functions for backward compatibility
TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "read_component_docs",
            "description": "Read documentation for a component category.",
            "parameters": {
                "type": "object",
                "properties": {
                    "category": {
                        "type": "string",
                        "enum": ["basic", "layout", "form", "data", "chart", "feedback"],
                        "description": "The component category"
                    }
                },
                "required": ["category"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "complete_task",
            "description": "Submit the final UI schema.",
            "parameters": {
                "type": "object",
                "properties": {
                    "schema": {
                        "type": "object",
                        "description": "The UI schema object"
                    }
                },
                "required": ["schema"]
            }
        }
    }
]


class GenerateUIRequest(BaseModel):
    prompt: str
    # Optional agent phases - can be disabled by user
    enable_data_binding: bool = True  # DataBindingAgent phase
    enable_ux_enhancement: bool = True  # UXEnhancementAgent phase


class ModifyUIRequest(BaseModel):
    prompt: str
    current_schema: dict


@app.get("/health")
async def health():
    return {"status": "ok"}


@app.post("/api/generate-ui/stream")
async def generate_ui_stream(request: GenerateUIRequest):
    """Stream UI generation with real-time JSON output.

    All requests use multi-agent hierarchical architecture:
    1. SkeletonAgent generates structure with $placeholder markers
    2. ExpansionAgents expand placeholders in parallel
    3. DataBindingAgent adds data bindings (optional)
    4. UXEnhancementAgent improves interactions (optional)
    """

    async def event_generator() -> AsyncGenerator[str, None]:
        try:
            # Always use hierarchical multi-agent generation
            async for event in run_hierarchical_generation(
                request.prompt,
                enable_data_binding=request.enable_data_binding,
                enable_ux_enhancement=request.enable_ux_enhancement
            ):
                event_type = event["event"]
                event_data = json.dumps(event["data"], ensure_ascii=False)
                yield f"event: {event_type}\ndata: {event_data}\n\n"

        except Exception as e:
            error_data = json.dumps({"error": str(e)})
            yield f"event: error\ndata: {error_data}\n\n"

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


@app.post("/api/modify-ui/stream")
async def modify_ui_stream(request: ModifyUIRequest):
    """Stream UI modification with real-time JSON output."""

    async def event_generator() -> AsyncGenerator[str, None]:
        try:
            async for event in run_streaming_modification(request.prompt, request.current_schema):
                event_type = event["event"]
                event_data = json.dumps(event["data"], ensure_ascii=False)
                yield f"event: {event_type}\ndata: {event_data}\n\n"

        except Exception as e:
            error_data = json.dumps({"error": str(e)})
            yield f"event: error\ndata: {error_data}\n\n"

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


# =============================================================================
# SESSION MANAGEMENT API
# =============================================================================

@app.get("/api/sessions")
async def list_sessions():
    """List all saved generation sessions."""
    sessions = SchemaStore.list_sessions()
    return {"sessions": sessions}


@app.get("/api/sessions/{session_id}")
async def get_session(session_id: str):
    """Get a specific session's data."""
    store = SchemaStore(session_id)
    meta = store.load_session_meta()
    if not meta:
        raise HTTPException(status_code=404, detail="Session not found")

    return {
        "session_id": session_id,
        "meta": meta,
        "skeleton": store.load_skeleton(),
        "merged": store.load_merged(),
        "fragments": store.list_fragments()
    }


@app.get("/api/sessions/{session_id}/schema")
async def get_session_schema(session_id: str):
    """Get the merged schema for a session."""
    store = SchemaStore(session_id)
    schema = store.load_merged()
    if not schema:
        raise HTTPException(status_code=404, detail="Schema not found")
    return {"schema": schema}


@app.get("/api/sessions/{session_id}/fragments")
async def get_session_fragments(session_id: str):
    """Get all fragments for a session."""
    store = SchemaStore(session_id)
    fragments = {}
    for fid in store.list_fragments():
        fragments[fid] = store.load_fragment(fid)
    return {"fragments": fragments}


@app.get("/api/sessions/{session_id}/fragments/{fragment_id}")
async def get_fragment(session_id: str, fragment_id: str):
    """Get a specific fragment."""
    store = SchemaStore(session_id)
    fragment = store.load_fragment(fragment_id)
    if not fragment:
        raise HTTPException(status_code=404, detail="Fragment not found")
    return fragment


@app.get("/api/sessions/{session_id}/skeleton-refs")
async def get_skeleton_with_refs(session_id: str):
    """Get the skeleton with $ref references (for incremental loading)."""
    store = SchemaStore(session_id)
    skeleton_refs = store.load_skeleton_with_refs()
    if not skeleton_refs:
        # Fallback: generate from skeleton
        skeleton = store.load_skeleton()
        if skeleton:
            skeleton_refs = store.save_skeleton_with_refs(skeleton)
    if not skeleton_refs:
        raise HTTPException(status_code=404, detail="Skeleton not found")
    return {"skeleton_refs": skeleton_refs}


@app.get("/api/sessions/{session_id}/context")
async def get_session_context(session_id: str, max_tokens: int = 2000):
    """Get compressed context for a session (for LLM context management)."""
    store = SchemaStore(session_id)
    meta = store.load_session_meta()
    if not meta:
        raise HTTPException(status_code=404, detail="Session not found")

    context = store.get_compressed_context(max_tokens)
    return {"context": context}


@app.get("/api/sessions/{session_id}/resolve")
async def resolve_session_refs(session_id: str):
    """Resolve all $ref references and return the complete schema."""
    store = SchemaStore(session_id)
    skeleton_refs = store.load_skeleton_with_refs()
    if not skeleton_refs:
        raise HTTPException(status_code=404, detail="Skeleton not found")

    resolved = store.resolve_refs(skeleton_refs)
    return {"schema": resolved}


@app.get("/api/sessions/{session_id}/fragments-by-depth/{depth}")
async def get_fragments_by_depth(session_id: str, depth: int):
    """Get all fragments at a specific depth level."""
    store = SchemaStore(session_id)
    fragments = store.get_fragments_by_depth(depth)
    return {"depth": depth, "fragments": fragments}


@app.delete("/api/sessions/{session_id}")
async def delete_session(session_id: str):
    """Delete a session and all its data."""
    try:
        SchemaStore.delete_session(session_id)
        return {"status": "deleted", "session_id": session_id}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


class ResumeSessionRequest(BaseModel):
    session_id: str
    continue_from: str = "expansion"  # "expansion", "data_binding", "ux_enhancement"


@app.post("/api/sessions/{session_id}/resume")
async def resume_session(session_id: str, request: ResumeSessionRequest):
    """Resume a session from a specific phase."""
    store = SchemaStore(session_id)
    meta = store.load_session_meta()
    if not meta:
        raise HTTPException(status_code=404, detail="Session not found")

    # Load the current state
    skeleton = store.load_skeleton()
    merged = store.load_merged() or skeleton

    if not merged:
        raise HTTPException(status_code=400, detail="No schema to resume from")

    # TODO: Implement resume logic based on continue_from phase
    return {
        "status": "resumed",
        "session_id": session_id,
        "continue_from": request.continue_from,
        "schema": merged
    }


def build_frontend():
    """Build frontend if not already built."""
    if not DIST_DIR.exists() or not (DIST_DIR / "index.html").exists():
        logger.info("Building frontend...")
        result = subprocess.run(
            ["npm", "run", "build"],
            cwd=FRONTEND_DIR,
            capture_output=True,
            text=True
        )
        if result.returncode != 0:
            logger.error(f"Frontend build failed: {result.stderr}")
            return False
        logger.info("Frontend built successfully")
    return True


def setup_static_files():
    """Mount static files for production mode."""
    from fastapi.responses import FileResponse
    
    # Serve static assets
    if (DIST_DIR / "assets").exists():
        app.mount("/assets", StaticFiles(directory=DIST_DIR / "assets"), name="assets")
    
    # Serve other static files
    for item in DIST_DIR.iterdir():
        if item.is_file() and item.name != "index.html":
            @app.get(f"/{item.name}")
            async def serve_static(filename=item.name):
                return FileResponse(DIST_DIR / filename)
    
    # Catch-all route for SPA (must be last)
    @app.get("/{full_path:path}")
    async def serve_spa(full_path: str):
        # Don't serve index.html for API routes
        if full_path.startswith("api/") or full_path.startswith("health"):
            raise HTTPException(status_code=404)
        return FileResponse(DIST_DIR / "index.html")


def main():
    """
    Start the RealtimeUI application.
    
    Usage:
        uv run main          # Production mode (serves built frontend)
        uv run main --dev    # Development mode (API only, run frontend separately)
    """
    import argparse
    parser = argparse.ArgumentParser(description="RealtimeUI Server")
    parser.add_argument("--dev", action="store_true", help="Run in development mode (API only)")
    parser.add_argument("--host", default="0.0.0.0", help="Host to bind to")
    parser.add_argument("--port", type=int, default=8000, help="Port to bind to")
    parser.add_argument("--build", action="store_true", help="Force rebuild frontend")
    args = parser.parse_args()
    
    if args.dev:
        # Development mode: API only with hot reload
        logger.info("Starting in DEVELOPMENT mode (API only)")
        logger.info("Run frontend separately: cd frontend && npm run dev")
        uvicorn.run("main:app", host=args.host, port=args.port, reload=True)
    else:
        # Production mode: serve frontend + API
        if args.build or not DIST_DIR.exists():
            if not build_frontend():
                logger.error("Cannot start without frontend. Try: cd frontend && npm install && npm run build")
                sys.exit(1)

        # Setup static file serving
        setup_static_files()

        logger.info(f"Starting RealtimeUI at http://localhost:{args.port}")
        logger.info("Frontend + API served from single process")
        uvicorn.run(app, host=args.host, port=args.port)


if __name__ == "__main__":
    main()