"""
Claude Code Backend — Agentic coding backend powered by NVIDIA NIM models.
Exposes an OpenAI-compatible /v1/chat/completions endpoint with built-in
tools for file operations and bash execution.

Architecture:
  Space 1 (better-chatbot) --> this backend --> NVIDIA NIM API
  
The agentic loop:
  1. Receive user message from Space 1
  2. Send to NIM model with tool definitions
  3. If model returns tool_calls, execute them and loop
  4. If model returns text, stream it back to Space 1
  5. Persist conversation in Postgres
"""

import os
import json
import uuid
import subprocess
import asyncio
import time
import re
import collections
from pathlib import Path
from typing import AsyncIterator, Optional

from fastapi import FastAPI, Request, Header, HTTPException
from fastapi.responses import StreamingResponse, JSONResponse, HTMLResponse
from fastapi.middleware.cors import CORSMiddleware
from openai import AsyncOpenAI
import anyio
import asyncpg

# ---------------------------------------------------------------------------
# Globals & Activity Logs
# ---------------------------------------------------------------------------
activity_logs = collections.deque(maxlen=100)
MODEL_STATUSES = {}
ACTIVE_SESSIONS = set()

def log_activity(msg: str):
    timestamp = time.strftime("%H:%M:%S")
    log_line = f"[{timestamp}] {msg}"
    activity_logs.append(log_line)
    print(log_line)

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

NIM_API_KEY = os.environ.get("NVIDIA_NIM_API_KEY", "")
BACKEND_API_KEY = os.environ.get("BACKEND_API_KEY", "")
DATABASE_URL = os.environ.get("DATABASE_URL", "")
WORKSPACE_DIR = os.environ.get("WORKSPACE_DIR", "/tmp/workspace")
MAX_TOOL_ROUNDS = int(os.environ.get("MAX_TOOL_ROUNDS", "10"))

# NIM models that reliably support tool/function calling
TOOL_CAPABLE_MODELS = {
    "nvidia/nemotron-3-ultra-550b-a55b": "Nemotron 3 Ultra 550B (Agentic)",
    "z-ai/glm-5.1": "GLM 5.1 (Agentic)",
    "moonshotai/kimi-k2.6": "Kimi K2.6 (Agentic)",
    "minimaxai/minimax-m3": "MiniMax M3 (Agentic)",
    "stepfun-ai/step-3.7-flash": "Step 3.7 Flash (Agentic)",
    "minimaxai/minimax-m2.7": "MiniMax M2.7 (Agentic)",
    "meta/llama-3.1-70b-instruct": "Llama 3.1 70B (Agentic)",
    "meta/llama-3.1-405b-instruct": "Llama 3.1 405B (Agentic)",
    "qwen/qwen2.5-coder-32b-instruct": "Qwen 2.5 Coder 32B (Agentic)",
    "nvidia/llama-3.1-nemotron-70b-instruct": "Nemotron 70B (Agentic)",
    "meta/llama-3.3-70b-instruct": "Llama 3.3 70B (Agentic)",
}

# All models (tool-capable get agentic mode, others get plain chat)
ALL_MODELS = {
    **TOOL_CAPABLE_MODELS,
    "deepseek-ai/deepseek-r1": "DeepSeek R1 (Chat only)",
    "mistralai/mistral-large-2-instruct": "Mistral Large 2 (Chat only)",
}

RECOMMENDED_MODEL = "nvidia/llama-3.1-nemotron-70b-instruct"

# Ensure workspace exists
Path(WORKSPACE_DIR).mkdir(parents=True, exist_ok=True)

# ---------------------------------------------------------------------------
# NIM Client
# ---------------------------------------------------------------------------

nim_client = AsyncOpenAI(
    base_url="https://integrate.api.nvidia.com/v1",
    api_key=NIM_API_KEY,
)

# ---------------------------------------------------------------------------
# Tool Definitions (OpenAI function calling format)
# ---------------------------------------------------------------------------

TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "read_file",
            "description": "Read the contents of a file. Use this to inspect existing code, configs, or any text file.",
            "parameters": {
                "type": "object",
                "properties": {
                    "path": {
                        "type": "string",
                        "description": "Relative path to the file from the workspace root"
                    }
                },
                "required": ["path"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "write_file",
            "description": "Write content to a file. Creates the file if it doesn't exist, overwrites if it does. Creates parent directories automatically.",
            "parameters": {
                "type": "object",
                "properties": {
                    "path": {
                        "type": "string",
                        "description": "Relative path to the file from the workspace root"
                    },
                    "content": {
                        "type": "string",
                        "description": "The full content to write to the file"
                    }
                },
                "required": ["path", "content"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "run_bash",
            "description": "Execute a bash command in the workspace directory. Use for installing packages, running scripts, git operations, etc. Commands run with a 30 second timeout.",
            "parameters": {
                "type": "object",
                "properties": {
                    "command": {
                        "type": "string",
                        "description": "The bash command to execute"
                    }
                },
                "required": ["command"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "list_directory",
            "description": "List files and directories in a given path. Shows file sizes and directory markers.",
            "parameters": {
                "type": "object",
                "properties": {
                    "path": {
                        "type": "string",
                        "description": "Relative path to the directory from workspace root. Use '.' for the workspace root."
                    }
                },
                "required": ["path"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "grep_search",
            "description": "Search for a pattern in files within the workspace. Returns matching lines with file paths and line numbers.",
            "parameters": {
                "type": "object",
                "properties": {
                    "pattern": {
                        "type": "string",
                        "description": "The search pattern (supports basic regex)"
                    },
                    "path": {
                        "type": "string",
                        "description": "Directory or file to search in, relative to workspace root. Defaults to '.'",
                    }
                },
                "required": ["pattern"]
            }
        }
    },
]

# ---------------------------------------------------------------------------
# Tool Execution
# ---------------------------------------------------------------------------

def _safe_path(rel_path: str) -> Path:
    """Resolve a relative path safely within the workspace."""
    workspace = Path(WORKSPACE_DIR).resolve()
    target = (workspace / rel_path).resolve()
    # Prevent path traversal
    if not str(target).startswith(str(workspace)):
        raise ValueError(f"Path traversal detected: {rel_path}")
    return target


def repair_arguments(func_name: str, args: dict) -> tuple[dict, list[str]]:
    notes = []
    repaired_args = dict(args)
    
    # 1. Nesting extraction (e.g. {"path": {"path": "file.txt"}})
    for key in list(repaired_args.keys()):
        val = repaired_args[key]
        if isinstance(val, dict) and key in val:
            repaired_args[key] = val[key]
            notes.append(f"Flattened nested parameter '{key}'")
            
    # 2. Markdown stripping from bash command
    if func_name == "run_bash" and "command" in repaired_args:
        cmd = repaired_args["command"]
        if isinstance(cmd, str):
            pattern = r"```(?:bash)?\s*(.*?)\s*```"
            match = re.search(pattern, cmd, re.DOTALL)
            if match:
                repaired_args["command"] = match.group(1).strip()
                notes.append("Stripped markdown code blocks from bash command")
                
    # 3. Stringified array conversion
    for key, val in repaired_args.items():
        if isinstance(val, str) and val.strip().startswith("[") and val.strip().endswith("]"):
            try:
                parsed_arr = json.loads(val)
                if isinstance(parsed_arr, list):
                    repaired_args[key] = parsed_arr
                    notes.append(f"Converted stringified array for parameter '{key}' to native array")
            except:
                pass
                
    # 4. Optional empty objects replacing Null
    for key in list(repaired_args.keys()):
        if repaired_args[key] == {}:
            repaired_args[key] = None
            notes.append(f"Replaced empty object for parameter '{key}' with null")
            
    return repaired_args, notes


def execute_tool(name: str, arguments: dict) -> str:
    """Execute a tool and return its output as a string."""
    try:
        if name == "read_file":
            path = _safe_path(arguments["path"])
            if not path.exists():
                return f"Error: File not found: {arguments['path']}"
            if not path.is_file():
                return f"Error: Not a file: {arguments['path']}"
            content = path.read_text(encoding="utf-8", errors="replace")
            if len(content) > 50000:
                return content[:50000] + f"\n\n[Truncated — file is {len(content)} chars]"
            return content

        elif name == "write_file":
            path = _safe_path(arguments["path"])
            path.parent.mkdir(parents=True, exist_ok=True)
            path.write_text(arguments["content"], encoding="utf-8")
            return f"Successfully wrote {len(arguments['content'])} chars to {arguments['path']}"

        elif name == "run_bash":
            command = arguments["command"]
            # Safety: block dangerous commands
            blocked = ["rm -rf /", "mkfs", "dd if=", ":(){", "fork bomb"]
            if any(b in command.lower() for b in blocked):
                return "Error: Command blocked for safety reasons"
            result = subprocess.run(
                ["bash", "-c", command],
                cwd=WORKSPACE_DIR,
                capture_output=True,
                text=True,
                timeout=30,
                env={**os.environ, "HOME": "/tmp", "PATH": os.environ.get("PATH", "/usr/local/bin:/usr/bin:/bin")},
            )
            output = ""
            if result.stdout:
                output += result.stdout
            if result.stderr:
                output += ("\n" if output else "") + f"[stderr] {result.stderr}"
            if result.returncode != 0:
                output += f"\n[exit code: {result.returncode}]"
            if not output:
                output = "[command completed with no output]"
            # Truncate very long outputs
            if len(output) > 20000:
                output = output[:20000] + f"\n\n[Truncated — output is {len(output)} chars]"
            return output

        elif name == "list_directory":
            path = _safe_path(arguments.get("path", "."))
            if not path.exists():
                return f"Error: Directory not found: {arguments.get('path', '.')}"
            if not path.is_dir():
                return f"Error: Not a directory: {arguments.get('path', '.')}"
            entries = []
            for item in sorted(path.iterdir()):
                if item.is_dir():
                    entries.append(f"  📁 {item.name}/")
                else:
                    size = item.stat().st_size
                    if size < 1024:
                        size_str = f"{size}B"
                    elif size < 1024 * 1024:
                        size_str = f"{size/1024:.1f}KB"
                    else:
                        size_str = f"{size/(1024*1024):.1f}MB"
                    entries.append(f"  📄 {item.name} ({size_str})")
            return f"Contents of {arguments.get('path', '.')}:\n" + "\n".join(entries) if entries else "Empty directory"

        elif name == "grep_search":
            pattern = arguments["pattern"]
            search_path = arguments.get("path", ".")
            path = _safe_path(search_path)
            result = subprocess.run(
                ["grep", "-rn", "--include=*", pattern, str(path)],
                capture_output=True,
                text=True,
                timeout=10,
                cwd=WORKSPACE_DIR,
            )
            output = result.stdout if result.stdout else "No matches found"
            if len(output) > 10000:
                output = output[:10000] + "\n\n[Truncated]"
            return output

        else:
            return f"Error: Unknown tool: {name}"

    except subprocess.TimeoutExpired:
        return "Error: Command timed out after 30 seconds"
    except ValueError as e:
        return f"Error: {str(e)}"
    except Exception as e:
        return f"Error executing {name}: {str(e)}"


# ---------------------------------------------------------------------------
# Database (Session Persistence)
# ---------------------------------------------------------------------------

db_pool: Optional[asyncpg.Pool] = None


async def init_db():
    """Initialize database connection pool and create tables."""
    global db_pool
    if not DATABASE_URL:
        return
    try:
        db_pool = await asyncpg.create_pool(
            DATABASE_URL, 
            ssl="require", 
            min_size=1, 
            max_size=3,
            max_inactive_connection_lifetime=300
        )
        async with db_pool.acquire() as conn:
            await conn.execute("""
                CREATE TABLE IF NOT EXISTS agent_sessions (
                    id BIGSERIAL PRIMARY KEY,
                    session_id TEXT NOT NULL,
                    role TEXT NOT NULL,
                    content TEXT,
                    tool_calls JSONB,
                    tool_call_id TEXT,
                    created_at TIMESTAMPTZ DEFAULT NOW()
                );
                CREATE INDEX IF NOT EXISTS idx_agent_sessions_sid ON agent_sessions(session_id);
            """)
    except Exception as e:
        print(f"[DB] Warning: Could not initialize database: {e}")
        db_pool = None


async def save_message(session_id: str, role: str, content: str = None, 
                       tool_calls: list = None, tool_call_id: str = None):
    """Save a message to the session store."""
    if not db_pool:
        return
    try:
        async with db_pool.acquire() as conn:
            await conn.execute(
                "INSERT INTO agent_sessions (session_id, role, content, tool_calls, tool_call_id) VALUES ($1, $2, $3, $4, $5)",
                session_id, role, content,
                json.dumps(tool_calls) if tool_calls else None,
                tool_call_id,
            )
    except Exception as e:
        print(f"[DB] Warning: Could not save message: {e}")


async def load_session(session_id: str) -> list:
    """Load conversation history from the session store."""
    if not db_pool:
        return []
    try:
        async with db_pool.acquire() as conn:
            rows = await conn.fetch(
                "SELECT role, content, tool_calls, tool_call_id FROM agent_sessions WHERE session_id = $1 ORDER BY id",
                session_id,
            )
            messages = []
            for row in rows:
                msg = {"role": row["role"]}
                if row["content"]:
                    msg["content"] = row["content"]
                if row["tool_calls"]:
                    msg["tool_calls"] = json.loads(row["tool_calls"])
                if row["tool_call_id"]:
                    msg["tool_call_id"] = row["tool_call_id"]
                messages.append(msg)
            return messages
    except Exception as e:
        print(f"[DB] Warning: Could not load session: {e}")
        return []


# ---------------------------------------------------------------------------
# SSE Chunk Formatting (OpenAI delta format)
# ---------------------------------------------------------------------------

def make_chunk(request_id: str, model: str, content: str = "", finish_reason: str = None) -> str:
    """Create an OpenAI-compatible SSE chunk."""
    delta = {}
    if content:
        delta["content"] = content
    if finish_reason and not content:
        delta = {}
    
    chunk = {
        "id": f"chatcmpl-{request_id}",
        "object": "chat.completion.chunk",
        "created": int(time.time()),
        "model": model,
        "choices": [{
            "index": 0,
            "delta": delta,
            "finish_reason": finish_reason,
        }],
    }
    return f"data: {json.dumps(chunk)}\n\n"


# ---------------------------------------------------------------------------
# FastAPI Application
# ---------------------------------------------------------------------------

app = FastAPI(title="Claude Code Backend", version="1.0.0")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


def auth(authorization: str = None):
    """Verify bearer token."""
    if not BACKEND_API_KEY:
        return  # No auth configured
    expected = f"Bearer {BACKEND_API_KEY}"
    if authorization != expected:
        raise HTTPException(status_code=401, detail="Unauthorized")


async def check_models_health():
    global RECOMMENDED_MODEL
    # Test only the unstable frontier models (the rest). 
    # The stable ones (Step 3.7 Flash, Nemotron 3 Ultra, Qwen 2.5 Coder) are always free/working.
    models_to_test = [
        "moonshotai/kimi-k2.6",
        "z-ai/glm-5.1",
        "minimaxai/minimax-m3",
        "minimaxai/minimax-m2.7",
        "meta/llama-3.1-405b-instruct",
    ]
    
    best_model = None
    best_latency = 999.0
    
    # Mark stable models as permanently ONLINE in the status map
    stable_models = [
        "stepfun-ai/step-3.7-flash",
        "nvidia/nemotron-3-ultra-550b-a55b",
        "qwen/qwen2.5-coder-32b-instruct"
    ]
    for model in stable_models:
        MODEL_STATUSES[model] = {"status": "ONLINE (Stable)", "latency": "Fast", "raw_latency": 0.1}

    log_activity("Periodic health check started: verifying unstable frontier NIM models...")
    for model in models_to_test:
        start_time = time.time()
        try:
            # Send a fast test prompt
            async with anyio.fail_after(15.0): # 15 seconds max timeout
                await nim_client.chat.completions.create(
                    model=model,
                    messages=[{"role": "user", "content": "1+1="}],
                    max_tokens=3,
                )
            latency = time.time() - start_time
            MODEL_STATUSES[model] = {"status": "ONLINE", "latency": f"{latency:.2f}s", "raw_latency": latency}
            log_activity(f"Model checked: {model} is ONLINE ({latency:.2f}s)")
            
            # Choose the fastest online unstable model
            if latency < best_latency:
                best_latency = latency
                best_model = model
                
        except Exception as e:
            MODEL_STATUSES[model] = {"status": "OFFLINE", "latency": "N/A", "raw_latency": 999.0}
            log_activity(f"Model checked: {model} is OFFLINE / TIMEOUT: {e}")
            
    if best_model:
        RECOMMENDED_MODEL = best_model
        log_activity(f"Best frontier model selected: {RECOMMENDED_MODEL} ({best_latency:.2f}s)")
    else:
        # Fallback to the stable Step 3.7 Flash if all frontier models are offline/throttled
        RECOMMENDED_MODEL = "stepfun-ai/step-3.7-flash"
        log_activity(f"All frontier models offline. Falling back to stable recommended model: {RECOMMENDED_MODEL}")

async def periodic_health_check_loop():
    # Wait 10 seconds after startup before the first check to let the space boot fully
    await asyncio.sleep(10)
    while True:
        try:
            await check_models_health()
        except Exception as e:
            log_activity(f"Health check loop error: {e}")
        await asyncio.sleep(900) # every 15 minutes (reduce frequency to save quota)

@app.on_event("startup")
async def startup():
    await init_db()
    Path(WORKSPACE_DIR).mkdir(parents=True, exist_ok=True)
    # Initialize statuses for all models
    for model_id, display_name in ALL_MODELS.items():
        MODEL_STATUSES[model_id] = {"status": "UNCHECKED", "latency": "N/A", "raw_latency": 999.0}
    # Start background health checking
    asyncio.create_task(periodic_health_check_loop())
    log_activity(f"FastAPI backend started. Workspace: {WORKSPACE_DIR}")


# ---------------------------------------------------------------------------
# /v1/chat/completions — Main endpoint
# ---------------------------------------------------------------------------

AGENTIC_SYSTEM_PROMPT = """You are an expert coding assistant with access to tools for file operations and command execution. 

When the user asks you to create, edit, or debug code:
1. Use `list_directory` and `read_file` to understand the current state
2. Use `write_file` to create or modify files
3. Use `run_bash` to execute commands (install packages, run scripts, test code)
4. Use `grep_search` to find patterns in code

IMPORTANT RULES:
- Always use tools to take action. Do NOT just describe what to do — actually DO it.
- After writing code, run it to verify it works.
- If a command fails, read the error and fix it.
- Work in the /tmp/workspace directory.
- Be concise in your explanations, but thorough in your tool usage.
"""


@app.post("/v1/chat/completions")
async def chat_completions(request: Request, authorization: str = Header(None)):
    auth(authorization)
    body = await request.json()

    requested_model = body.get("model", "meta/llama-3.1-70b-instruct")
    messages = body.get("messages", [])
    stream = body.get("stream", False)
    session_id = body.get("session_id") or str(uuid.uuid4())
    
    is_agentic = requested_model in TOOL_CAPABLE_MODELS
    request_id = str(uuid.uuid4())[:8]

    ACTIVE_SESSIONS.add(session_id)
    log_activity(f"Session [{session_id[:6]}] connected. Model: {requested_model}")

    # Build message history
    final_messages = []
    
    # Add agentic system prompt for tool-capable models
    if is_agentic:
        # Check if there's already a system message
        has_system = any(m.get("role") == "system" for m in messages)
        if has_system:
            # Prepend agentic prompt to existing system message
            for m in messages:
                if m["role"] == "system":
                    final_messages.append({
                        "role": "system",
                        "content": AGENTIC_SYSTEM_PROMPT + "\n\nAdditional instructions:\n" + m["content"]
                    })
                else:
                    final_messages.append(m)
        else:
            final_messages.append({"role": "system", "content": AGENTIC_SYSTEM_PROMPT})
            final_messages.extend(messages)
    else:
        final_messages = list(messages)

    # Save the user's message to DB
    user_msg = next((m for m in reversed(messages) if m.get("role") == "user"), None)
    if user_msg:
        await save_message(session_id, "user", user_msg.get("content", ""))

    if not stream:
        # Non-streaming: simple completion
        try:
            kwargs = {"model": requested_model, "messages": final_messages}
            if is_agentic:
                kwargs["tools"] = TOOLS
                kwargs["tool_choice"] = "auto"
            response = await nim_client.chat.completions.create(**kwargs)
            content = response.choices[0].message.content or ""
            await save_message(session_id, "assistant", content)
            ACTIVE_SESSIONS.discard(session_id)
            log_activity(f"Session [{session_id[:6]}] finished (non-streaming)")
            return JSONResponse({
                "id": f"chatcmpl-{request_id}",
                "object": "chat.completion",
                "created": int(time.time()),
                "model": requested_model,
                "choices": [{"index": 0, "message": {"role": "assistant", "content": content}, "finish_reason": "stop"}],
            })
        except Exception as e:
            ACTIVE_SESSIONS.discard(session_id)
            return JSONResponse({"error": {"message": str(e), "type": "internal_error"}}, status_code=500)

    # Streaming + agentic loop
    async def generate() -> AsyncIterator[str]:
        nonlocal final_messages
        try:
            for round_num in range(MAX_TOOL_ROUNDS + 1):
                kwargs = {"model": requested_model, "messages": final_messages, "stream": True}
                if is_agentic:
                    kwargs["tools"] = TOOLS
                    kwargs["tool_choice"] = "auto"

                # Collect streamed response
                full_content = ""
                tool_calls_raw = {}  # index -> {id, name, arguments_str}
                
                async for chunk in await nim_client.chat.completions.create(**kwargs):
                    choice = chunk.choices[0] if chunk.choices else None
                    if not choice:
                        continue
                    delta = choice.delta
                    
                    # Stream text content to client
                    if delta and delta.content:
                        full_content += delta.content
                        yield make_chunk(request_id, requested_model, delta.content)
                    
                    # Collect tool calls
                    if delta and delta.tool_calls:
                        for tc in delta.tool_calls:
                            idx = tc.index
                            if idx not in tool_calls_raw:
                                tool_calls_raw[idx] = {
                                    "id": tc.id or f"call_{uuid.uuid4().hex[:8]}",
                                    "name": tc.function.name if tc.function and tc.function.name else "",
                                    "arguments": ""
                                }
                            if tc.function and tc.function.name:
                                tool_calls_raw[idx]["name"] = tc.function.name
                            if tc.id:
                                tool_calls_raw[idx]["id"] = tc.id
                            if tc.function and tc.function.arguments:
                                tool_calls_raw[idx]["arguments"] += tc.function.arguments
                    
                    # Check for finish
                    if choice.finish_reason == "stop":
                        break
                    if choice.finish_reason == "tool_calls":
                        break

                # If no tool calls, we're done
                if not tool_calls_raw:
                    await save_message(session_id, "assistant", full_content)
                    yield make_chunk(request_id, requested_model, finish_reason="stop")
                    yield "data: [DONE]\n\n"
                    return

                # Execute tool calls
                tool_calls_list = []
                for idx in sorted(tool_calls_raw.keys()):
                    tc = tool_calls_raw[idx]
                    tool_calls_list.append({
                        "id": tc["id"],
                        "type": "function",
                        "function": {"name": tc["name"], "arguments": tc["arguments"]}
                    })

                # Add assistant message with tool calls to history
                assistant_msg = {"role": "assistant", "content": full_content or None, "tool_calls": tool_calls_list}
                final_messages.append(assistant_msg)

                # Execute each tool and add results
                for tc in tool_calls_list:
                    func_name = tc["function"]["name"]
                    raw_args_str = tc["function"]["arguments"]
                    try:
                        func_args = json.loads(raw_args_str)
                    except json.JSONDecodeError:
                        # Attempt raw JSON repair
                        repaired_str = raw_args_str.strip()
                        if not repaired_str.startswith("{"):
                            repaired_str = "{" + repaired_str
                        if not repaired_str.endswith("}"):
                            repaired_str = repaired_str + "}"
                        try:
                            func_args = json.loads(repaired_str)
                            log_activity(f"Auto-fixed invalid JSON string for tool: {func_name}")
                        except:
                            func_args = {}

                    # Perform semantic repairs
                    repaired_args, repair_notes = repair_arguments(func_name, func_args)
                    
                    # Log activity
                    log_activity(f"Tool execution: {func_name} args={repaired_args}")
                    if repair_notes:
                        for note in repair_notes:
                            log_activity(f"[Tool Repair] {note}")

                    # Show tool execution to user
                    yield make_chunk(request_id, requested_model, f"\n\n🔧 **{func_name}**")
                    if repair_notes:
                        yield make_chunk(request_id, requested_model, " *(Auto-Repaired)*")
                    
                    if func_name == "run_bash" and "command" in repaired_args:
                        yield make_chunk(request_id, requested_model, f": `{repaired_args['command']}`\n")
                    elif func_name == "read_file" and "path" in repaired_args:
                        yield make_chunk(request_id, requested_model, f": `{repaired_args['path']}`\n")
                    elif func_name == "write_file" and "path" in repaired_args:
                        yield make_chunk(request_id, requested_model, f": `{repaired_args['path']}`\n")
                    elif func_name == "list_directory":
                        yield make_chunk(request_id, requested_model, f": `{repaired_args.get('path', '.')}`\n")
                    elif func_name == "grep_search":
                        yield make_chunk(request_id, requested_model, f": `{repaired_args.get('pattern', '')}`\n")
                    else:
                        yield make_chunk(request_id, requested_model, "\n")

                    # Execute the tool
                    result = execute_tool(func_name, repaired_args)

                    # Append teaching note if repaired
                    if repair_notes:
                        result += f"\n\n[SYSTEM REPAIR NOTE: The harness automatically fixed formatting issues: {', '.join(repair_notes)}. Please strictly follow the tool's JSON schema in subsequent calls without these wrapping/formatting errors.]"

                    # Show truncated result to user
                    preview = result[:500] + ("..." if len(result) > 500 else "")
                    yield make_chunk(request_id, requested_model, f"```\n{preview}\n```\n")

                    # Add tool result to message history
                    final_messages.append({
                        "role": "tool",
                        "tool_call_id": tc["id"],
                        "content": result,
                    })

                    await save_message(session_id, "tool", result, tool_call_id=tc["id"])

                # Continue the agentic loop (model processes tool results)

            # If we hit max rounds, finish
            yield make_chunk(request_id, requested_model, "\n\n⚠️ Reached maximum tool call rounds.")
            yield make_chunk(request_id, requested_model, finish_reason="stop")
            yield "data: [DONE]\n\n"

        except Exception as e:
            error_msg = f"\n\n❌ Error: {str(e)}"
            yield make_chunk(request_id, requested_model, error_msg)
            yield make_chunk(request_id, requested_model, finish_reason="stop")
            yield "data: [DONE]\n\n"

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


# ---------------------------------------------------------------------------
# /v1/models — Model listing
# ---------------------------------------------------------------------------

@app.get("/v1/models")
async def list_models(authorization: str = Header(None)):
    auth(authorization)
    models = []
    for model_id, display_name in ALL_MODELS.items():
        models.append({
            "id": model_id,
            "object": "model",
            "created": 1700000000,
            "owned_by": "nvidia-nim",
            "permission": [],
            "root": model_id,
            "parent": None,
        })
    return {"object": "list", "data": models}


# ---------------------------------------------------------------------------
# /health — Health check
# ---------------------------------------------------------------------------

@app.get("/api/workspace/tree")
async def get_workspace_tree():
    def build_tree(current_path: Path, relative_to: Path) -> dict:
        name = current_path.name
        try:
            rel_path = str(current_path.relative_to(relative_to)).replace("\\", "/")
        except ValueError:
            rel_path = ""
        if rel_path == ".":
            rel_path = ""
            
        if current_path.is_dir():
            children = []
            try:
                for child in sorted(current_path.iterdir(), key=lambda x: (not x.is_dir(), x.name)):
                    if child.name in [".git", "node_modules", ".next", "__pycache__", ".agents", ".gemini"]:
                        continue
                    children.append(build_tree(child, relative_to))
            except Exception:
                pass
            return {
                "name": name or "workspace",
                "path": rel_path,
                "type": "directory",
                "children": children
            }
        else:
            return {
                "name": name,
                "path": rel_path,
                "type": "file",
                "size": current_path.stat().st_size if current_path.exists() else 0
            }
            
    try:
        w_path = Path(WORKSPACE_DIR).resolve()
        if not w_path.exists():
            w_path.mkdir(parents=True, exist_ok=True)
        return build_tree(w_path, w_path)
    except Exception as e:
        return {"error": str(e)}


@app.get("/api/workspace/file")
async def get_workspace_file(path: str):
    try:
        safe_p = _safe_path(path)
        if not safe_p.exists() or not safe_p.is_file():
            raise HTTPException(status_code=404, detail="File not found")
        content = safe_p.read_text(encoding="utf-8", errors="replace")
        return {"path": path, "content": content}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# ---------------------------------------------------------------------------
# Dashboard and Status API
# ---------------------------------------------------------------------------

DASHBOARD_HTML = """
<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Claude Code Agent Console</title>
    <script src="https://cdn.tailwindcss.com"></script>
    <style>
        @import url('https://fonts.googleapis.com/css2?family=Fira+Code:wght@400;500;700&family=Outfit:wght@400;600;800&display=swap');
        body {
            font-family: 'Outfit', sans-serif;
            background-color: #0b0c10;
        }
        .code-font {
            font-family: 'Fira Code', monospace;
        }
        .glow-amber {
            box-shadow: 0 0 15px rgba(245, 158, 11, 0.2);
        }
    </style>
</head>
<body class="text-gray-100 min-h-screen flex flex-col pb-10">
    <header class="border-b border-gray-800 bg-gray-950/80 backdrop-blur px-6 py-4 flex items-center justify-between sticky top-0 z-50">
        <div class="flex items-center space-x-3">
            <span class="text-2xl font-extrabold tracking-tight bg-gradient-to-r from-blue-400 via-indigo-400 to-purple-400 bg-clip-text text-transparent">
                Claude Code Agent Console
            </span>
            <span class="px-2 py-0.5 text-xs rounded bg-blue-500/10 text-blue-400 border border-blue-500/20 font-semibold animate-pulse">
                LIVE
            </span>
        </div>
        <div class="flex items-center space-x-4 text-sm text-gray-400">
            <div>Workspace: <span class="text-gray-200 code-font">/tmp/workspace</span></div>
            <div class="h-4 w-px bg-gray-800"></div>
            <div>Active Sessions: <span id="active-sessions-count" class="text-blue-400 font-bold code-font">0</span></div>
        </div>
    </header>

    <!-- Navigation Tabs -->
    <div class="border-b border-gray-800 max-w-7xl w-full mx-auto px-6 mt-6 flex space-x-6 text-sm">
        <button onclick="switchTab('models')" id="tab-btn-models" class="pb-3 border-b-2 border-blue-500 font-semibold text-blue-400 transition-all">NIM Models</button>
        <button onclick="switchTab('logs')" id="tab-btn-logs" class="pb-3 border-b-2 border-transparent text-gray-400 hover:text-gray-200 font-semibold transition-all">Live Logs</button>
        <button onclick="switchTab('explorer')" id="tab-btn-explorer" class="pb-3 border-b-2 border-transparent text-gray-400 hover:text-gray-200 font-semibold flex items-center space-x-1 transition-all">
            <span>Workspace Explorer (IDE)</span>
            <span class="px-1.5 py-0.5 rounded bg-blue-500/10 text-blue-400 border border-blue-500/20 text-[10px] font-bold">VS Code View</span>
        </button>
    </div>

    <!-- MAIN SECTIONS -->
    <main class="max-w-7xl w-full mx-auto px-6 mt-8 flex-1">
        
        <!-- SECTION: Models -->
        <div id="section-models" class="space-y-6">
            <div class="flex items-center justify-between">
                <h2 class="text-lg font-bold tracking-tight text-gray-300">Nvidia NIM Models & Health Status</h2>
                <span class="text-xs text-gray-500">Checked every 15 mins</span>
            </div>
            <div id="models-container" class="grid grid-cols-1 md:grid-cols-2 lg:grid-cols-3 gap-4">
                <!-- Dynamically loaded models go here -->
            </div>
        </div>

        <!-- SECTION: Logs -->
        <div id="section-logs" class="hidden space-y-6">
            <h2 class="text-lg font-bold tracking-tight text-gray-300">System Activity Logs</h2>
            <div class="border border-gray-800 rounded-lg overflow-hidden bg-gray-950 flex flex-col min-h-[500px]">
                <div class="bg-gray-900 px-4 py-2 border-b border-gray-800 flex items-center justify-between">
                    <span class="text-xs text-gray-400 font-semibold code-font">agent-stdout.log</span>
                    <div class="flex space-x-1.5">
                        <span class="w-2.5 h-2.5 rounded-full bg-red-500/30"></span>
                        <span class="w-2.5 h-2.5 rounded-full bg-yellow-500/30"></span>
                        <span class="w-2.5 h-2.5 rounded-full bg-green-500/30"></span>
                    </div>
                </div>
                <div id="terminal-content" class="p-4 flex-1 overflow-y-auto code-font text-xs text-green-400 bg-black/90 space-y-1 select-all h-[450px]">
                    <!-- Logs go here -->
                </div>
            </div>
        </div>

        <!-- SECTION: Workspace Explorer -->
        <div id="section-explorer" class="hidden space-y-6">
            <div class="flex items-center justify-between">
                <h2 class="text-lg font-bold tracking-tight text-gray-300">Visual Workspace IDE</h2>
                <button onclick="refreshFileTree()" class="text-xs px-2.5 py-1 rounded bg-blue-500/10 text-blue-400 border border-blue-500/20 hover:bg-blue-500/20 transition-all font-semibold">
                    🔄 Refresh Tree
                </button>
            </div>
            
            <div class="grid grid-cols-1 md:grid-cols-3 gap-6 border border-gray-800 rounded-xl bg-gray-950 overflow-hidden h-[600px]">
                <!-- File Tree Sidebar -->
                <div class="border-r border-gray-800 flex flex-col bg-gray-950 h-full">
                    <div class="px-4 py-2 border-b border-gray-800 bg-gray-900 text-xs font-semibold tracking-wider text-gray-400 code-font">
                        📁 EXPLORER: WORKSPACE
                    </div>
                    <div id="file-tree" class="p-3 flex-1 overflow-y-auto space-y-0.5 select-none">
                        <!-- Tree will be loaded here -->
                        <span class="text-xs text-gray-500 italic px-2">Loading directory tree...</span>
                    </div>
                </div>

                <!-- Editor panel -->
                <div class="md:col-span-2 flex flex-col bg-black/40 h-full">
                    <div class="px-4 py-2 border-b border-gray-800 bg-gray-900 flex items-center justify-between">
                        <span id="editor-title" class="text-xs font-semibold text-gray-400 code-font">📄 Welcome screen</span>
                        <div class="flex space-x-1.5">
                            <span class="w-2 h-2 rounded-full bg-gray-700"></span>
                            <span class="w-2 h-2 rounded-full bg-gray-700"></span>
                        </div>
                    </div>
                    <div class="flex-1 p-4 overflow-auto code-font text-xs text-gray-200">
                        <pre id="editor-content" class="whitespace-pre overflow-x-auto select-text h-[500px]">
Welcome to Claude Code Workspace Explorer.
Select a file from the sidebar explorer on the left to read its code contents in real-time.
                        </pre>
                    </div>
                </div>
            </div>
        </div>

    </main>

    <script>
        let currentTab = 'models';

        function switchTab(tabId) {
            currentTab = tabId;
            // Toggle sections
            document.getElementById('section-models').classList.add('hidden');
            document.getElementById('section-logs').classList.add('hidden');
            document.getElementById('section-explorer').classList.add('hidden');
            
            document.getElementById('tab-btn-models').className = 'pb-3 border-b-2 border-transparent text-gray-400 hover:text-gray-200 font-semibold transition-all';
            document.getElementById('tab-btn-logs').className = 'pb-3 border-b-2 border-transparent text-gray-400 hover:text-gray-200 font-semibold transition-all';
            document.getElementById('tab-btn-explorer').className = 'pb-3 border-b-2 border-transparent text-gray-400 hover:text-gray-200 font-semibold flex items-center space-x-1 transition-all';

            if (tabId === 'models') {
                document.getElementById('section-models').classList.remove('hidden');
                document.getElementById('tab-btn-models').className = 'pb-3 border-b-2 border-blue-500 font-semibold text-blue-400 transition-all';
            } else if (tabId === 'logs') {
                document.getElementById('section-logs').classList.remove('hidden');
                document.getElementById('tab-btn-logs').className = 'pb-3 border-b-2 border-blue-500 font-semibold text-blue-400 transition-all';
            } else if (tabId === 'explorer') {
                document.getElementById('section-explorer').classList.remove('hidden');
                document.getElementById('tab-btn-explorer').className = 'pb-3 border-b-2 border-blue-500 font-semibold text-blue-400 flex items-center space-x-1 transition-all';
                refreshFileTree();
            }
        }

        async function fetchSystemData() {
            try {
                const res = await fetch('/health');
                if (!res.ok) return;
                const data = await res.json();
                document.getElementById('active-sessions-count').innerText = data.active_sessions || 0;
            } catch (e) {
                console.error(e);
            }
        }

        async function fetchModels() {
            try {
                const res = await fetch('/api/models-status');
                if (!res.ok) return;
                const models = await res.json();
                
                const container = document.getElementById('models-container');
                container.innerHTML = '';
                
                models.forEach(model => {
                    const isRec = model.is_recommended;
                    const isOnline = model.status.includes('ONLINE');
                    
                    const card = document.createElement('div');
                    card.className = `p-4 border rounded-xl bg-gray-950 transition-all ${
                        isRec ? 'border-amber-500/50 glow-amber bg-amber-500/5' : 'border-gray-800 bg-gray-950'
                    }`;
                    
                    card.innerHTML = `
                        <div class="flex items-center justify-between mb-3">
                            <span class="text-xs text-gray-500 code-font truncate max-w-[200px]" title="${model.id}">${model.id}</span>
                            <div class="flex items-center space-x-2">
                                ${isRec ? '<span class="text-[10px] px-1.5 py-0.5 rounded bg-amber-500/10 text-amber-400 border border-amber-500/20 font-bold">★ Recommended</span>' : ''}
                                <span class="h-2 w-2 rounded-full ${isOnline ? 'bg-green-500 animate-pulse' : 'bg-red-500'}"></span>
                                <span class="text-[10px] font-bold ${isOnline ? 'text-green-400' : 'text-red-400'}">${model.status}</span>
                            </div>
                        </div>
                        <h3 class="text-sm font-bold text-gray-200 mb-2 truncate">${model.name}</h3>
                        <div class="flex items-center justify-between text-xs text-gray-400 border-t border-gray-900 pt-2">
                            <span>Type: <strong class="text-gray-300 font-medium">${model.type}</strong></span>
                            <span>Latency: <strong class="text-blue-400 code-font">${model.latency}</strong></span>
                        </div>
                    `;
                    container.appendChild(card);
                });
            } catch (e) {
                console.error(e);
            }
        }

        async function fetchLogs() {
            if (currentTab !== 'logs') return;
            try {
                const res = await fetch('/api/logs');
                if (!res.ok) return;
                const logs = await res.json();
                
                const term = document.getElementById('terminal-content');
                const shouldScroll = term.scrollHeight - term.clientHeight <= term.scrollTop + 50;
                term.innerHTML = logs.map(line => `<div>${line}</div>`).join('');
                if (shouldScroll) {
                    term.scrollTop = term.scrollHeight;
                }
            } catch (e) {
                console.error(e);
            }
        }

        // File Explorer Logic
        async function refreshFileTree() {
            try {
                const res = await fetch('/api/workspace/tree');
                if (!res.ok) return;
                const root = await res.json();
                
                const container = document.getElementById('file-tree');
                container.innerHTML = renderNode(root);
            } catch (e) {
                console.error(e);
            }
        }

        function renderNode(node, depth = 0) {
            const isDir = node.type === 'directory';
            const icon = isDir ? '📁' : '📄';
            const indent = depth * 12;
            
            let html = `
                <div class="flex items-center py-1 px-2 hover:bg-gray-800 rounded cursor-pointer transition-all text-xs" 
                     style="padding-left: ${indent}px"
                     onclick="${isDir ? `toggleDir('${node.path}')` : `openFile('${node.path}')`}">
                    <span class="mr-2">${icon}</span>
                    <span class="truncate ${isDir ? 'text-gray-300 font-medium' : 'text-gray-400'}">${node.name}</span>
                </div>
            `;
            
            if (isDir && node.children && node.children.length > 0) {
                html += `<div id="dir-${node.path.replace(/\\/g, '-').replace(/\\//g, '-')}" class="space-y-0.5">`;
                node.children.forEach(child => {
                    html += renderNode(child, depth + 1);
                });
                html += `</div>`;
            } else if (isDir && (!node.children || node.children.length === 0)) {
                html += `<div class="text-[10px] text-gray-600 italic" style="padding-left: ${indent + 16}px">(empty)</div>`;
            }
            
            return html;
        }

        async function openFile(path) {
            document.getElementById('editor-title').innerText = `📄 ${path}`;
            document.getElementById('editor-content').innerText = "Loading file content...";
            try {
                const res = await fetch(`/api/workspace/file?path=${encodeURIComponent(path)}`);
                if (!res.ok) {
                    document.getElementById('editor-content').innerText = "Error: Failed to fetch file content.";
                    return;
                }
                const data = await res.json();
                document.getElementById('editor-content').innerText = data.content;
            } catch (e) {
                document.getElementById('editor-content').innerText = `Error: ${e.message}`;
            }
        }

        function toggleDir(path) {
            const safeId = `dir-${path.replace(/\\/g, '-').replace(/\\//g, '-')}`;
            const elem = document.getElementById(safeId);
            if (elem) {
                elem.classList.toggle('hidden');
            }
        }

        setInterval(fetchSystemData, 3000);
        setInterval(fetchModels, 3000);
        setInterval(fetchLogs, 2000);
        
        fetchSystemData();
        fetchModels();
        fetchLogs();
    </script>
</body>
</html>
"""

@app.get("/", response_class=HTMLResponse)
async def dashboard():
    return HTMLResponse(content=DASHBOARD_HTML)


@app.get("/api/logs")
async def get_logs():
    return list(activity_logs)


@app.get("/api/models-status")
async def get_models_status():
    status_list = []
    for model_id, display_name in ALL_MODELS.items():
        status_info = MODEL_STATUSES.get(model_id, {"status": "ONLINE (Unchecked)", "latency": "N/A"})
        is_rec = model_id == RECOMMENDED_MODEL
        is_agentic = model_id in TOOL_CAPABLE_MODELS
        status_list.append({
            "id": model_id,
            "name": display_name,
            "status": status_info["status"],
            "latency": status_info["latency"],
            "is_recommended": is_rec,
            "type": "Agentic (Tools)" if is_agentic else "Chat Only",
        })
    # Sort: Recommended first, then Agentic, then Chat
    status_list.sort(key=lambda m: (not m["is_recommended"], m["type"] != "Agentic (Tools)", m["name"]))
    return status_list


import shutil
import threading
import signal
from fastapi.responses import FileResponse

@app.get("/api/backup/download")
async def download_backup(authorization: str = Header(None)):
    auth(authorization)
    snapshot_dir = "/tmp/workspace_snapshot"
    archive_base = "/tmp/workspace_backup_download"
    archive_zip = archive_base + ".zip"
    
    # Clean up old files/folders
    for path in [snapshot_dir, archive_zip]:
        if os.path.exists(path):
            try:
                if os.path.isdir(path):
                    shutil.rmtree(path)
                else:
                    os.unlink(path)
            except Exception:
                pass
                
    try:
        # 1. Atomic-like snapshot copy (ignoring temporary files)
        shutil.copytree(WORKSPACE_DIR, snapshot_dir, symlinks=True, ignore=shutil.ignore_patterns('.git', 'node_modules', '.next'))
        # 2. Archive the snapshot folder to disk to prevent OOM memory spike
        shutil.make_archive(archive_base, 'zip', snapshot_dir)
        # 3. Clean up the snapshot directory immediately
        shutil.rmtree(snapshot_dir)
        
        if not os.path.exists(archive_zip):
            raise HTTPException(status_code=500, detail="Failed to create zip archive")
        return FileResponse(archive_zip, media_type="application/zip", filename="workspace_backup.zip")
    except Exception as e:
        if os.path.exists(snapshot_dir):
            shutil.rmtree(snapshot_dir)
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/health")
async def health():
    return {
        "status": "ok",
        "workspace": WORKSPACE_DIR,
        "workspace_exists": Path(WORKSPACE_DIR).exists(),
        "db_connected": db_pool is not None,
        "models_count": len(ALL_MODELS),
        "recommended_model": RECOMMENDED_MODEL,
        "active_sessions": len(ACTIVE_SESSIONS),
    }


# ---------------------------------------------------------------------------
# Watchdog Daemon for Claude Code Subprocesses (Orphan Reaper)
# ---------------------------------------------------------------------------

def run_watchdog():
    log_activity("System Watchdog Daemon started (PPID-based Orphan detection)")
    while True:
        try:
            import psutil
            for proc in psutil.process_iter(['pid', 'ppid', 'name', 'cmdline', 'status']):
                try:
                    cmd = " ".join(proc.info['cmdline'] or [])
                    # Match the CLI binary (looks like claude-code or anthropic CLI wrapper)
                    if "claude" in cmd.lower() or "anthropic" in cmd.lower():
                        ppid = proc.info['ppid']
                        pid = proc.info['pid']
                        
                        # Is the parent still alive and not a zombie?
                        parent_exists = False
                        if ppid != 1:  # Orphaned processes get reparented to PID 1 in Linux
                            try:
                                parent_proc = psutil.Process(ppid)
                                if parent_proc.is_running() and parent_proc.status() != psutil.STATUS_ZOMBIE:
                                    parent_exists = True
                            except psutil.NoSuchProcess:
                                pass
                        
                        if not parent_exists:
                            log_activity(f"[Watchdog SIGKILL] Reaping orphaned Claude Code process PID {pid} (PPID {ppid})")
                            proc.terminate()
                            time.sleep(2)
                            if proc.is_running():
                                proc.kill()
                except Exception:
                    continue
        except ImportError:
            # Fallback zero-dependency shell parser using /proc
            try:
                # Find all processes and examine their parent PID
                out = subprocess.check_output("ps -o pid,ppid,args | grep -E 'claude|anthropic' | grep -v grep", shell=True, text=True)
                for line in out.strip().split("\n"):
                    parts = line.strip().split(None, 2)
                    if len(parts) >= 2:
                        pid = int(parts[0])
                        ppid = int(parts[1])
                        
                        # Check if parent pid exists/is alive
                        parent_exists = False
                        if ppid != 1:
                            # Check /proc/[ppid] directory
                            if os.path.exists(f"/proc/{ppid}"):
                                parent_exists = True
                                
                        if not parent_exists:
                            log_activity(f"[Watchdog SIGKILL Fallback] Reaping orphaned process PID {pid} (PPID {ppid})")
                            try:
                                os.kill(pid, signal.SIGTERM)
                                time.sleep(2)
                                os.kill(pid, signal.SIGKILL)
                            except Exception:
                                pass
            except Exception:
                pass
        except Exception as e:
            log_activity(f"[Watchdog Error] {e}")
        time.sleep(60)


async def db_heartbeat_loop():
    log_activity("Database Heartbeat task started")
    while True:
        try:
            if db_pool:
                async with db_pool.acquire() as conn:
                    await conn.execute("SELECT 1")
                    log_activity("[Heartbeat] Pinged Aiven PostgreSQL successfully")
        except Exception as e:
            log_activity(f"[Heartbeat Warning] Failed to ping database: {e}")
        await asyncio.sleep(240) # Every 4 minutes


@app.on_event("startup")
async def startup_event():
    # Start the watchdog thread on startup
    threading.Thread(target=run_watchdog, daemon=True).start()
    # Start the db keep-alive loop on FastAPI event loop
    asyncio.create_task(db_heartbeat_loop())


# ---------------------------------------------------------------------------
# Entrypoint
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)