app.py

import os
import re
import json
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Any

import gradio as gr


# ----------------------
# Data loading utilities
# ----------------------

HERE = Path(__file__).parent
DATASOURCE_TXT = HERE / "datasource.txt"
STATIC_DATA_DIR = HERE / "static_data"


def _slugify(text: str) -> str:
    text = text.strip().lower()
    text = re.sub(r"[^a-z0-9\s_-]+", "", text)
    text = re.sub(r"[\s_-]+", "-", text).strip("-")
    return text or "uncategorized"


def _titleize_slug(slug: str) -> str:
    return re.sub(r"[-_]+", " ", slug).title()


def _read_text(path: Path) -> Optional[str]:
    try:
        return path.read_text(encoding="utf-8")
    except Exception:
        return None


def _try_parse_json(text: str) -> Optional[dict]:
    try:
        return json.loads(text)
    except Exception:
        return None


def _try_parse_yaml(text: str) -> Optional[dict]:
    try:
        import yaml  # type: ignore

        return yaml.safe_load(text)
    except Exception:
        return None


def _extract_field(obj: dict, keys: List[str]) -> Optional[str]:
    for k in keys:
        if k in obj and isinstance(obj[k], str) and obj[k].strip():
            return obj[k].strip()
    # common nested forms, like { system: { content: "..." } } or messages: [{role: system, content: "..."}]
    if "system" in obj and isinstance(obj["system"], dict):
        for k in ("content", "prompt", "text"):
            v = obj["system"].get(k)
            if isinstance(v, str) and v.strip():
                return v.strip()
    if "messages" in obj and isinstance(obj["messages"], list):
        for m in obj["messages"]:
            if isinstance(m, dict) and m.get("role") == "system":
                content = m.get("content")
                if isinstance(content, str) and content.strip():
                    return content.strip()
    return None


def _extract_agent_from_obj(obj: dict, fallback_category: str, source_path: Path) -> Optional[dict]:
    # Heuristics to recognize agent-like configs.
    raw_name = _extract_field(obj, ["name", "agent_name", "title", "id"])
    description = _extract_field(obj, ["description", "desc", "about", "summary"])
    system_prompt = _extract_field(
        obj,
        [
            "system_prompt",
            "prompt",
            "instructions",
            "system",
            "system_instructions",
            "system_text",
        ],
    )

    if not (raw_name and system_prompt):
        return None

    # Normalize the name for display
    name = _titleize_slug(_slugify(raw_name))
    
    category = (
        _extract_field(obj, ["category", "group", "type"]) or fallback_category or "uncategorized"
    )
    category_slug = _slugify(category)
    agent_id = _slugify(f"{category_slug}-{raw_name}")

    return {
        "id": agent_id,
        "name": name,
        "description": description or "",
        "system_prompt": system_prompt,
        "category": category_slug,
        "source": str(source_path.relative_to(HERE)),
    }


def _parse_markdown_frontmatter(text: str) -> Optional[dict]:
    """Parse YAML frontmatter from markdown files and include body content."""
    if not text.startswith('---'):
        return None
    
    # Find the end of frontmatter
    lines = text.split('\n')
    end_idx = -1
    for i, line in enumerate(lines[1:], 1):
        if line.strip() == '---':
            end_idx = i
            break
    
    if end_idx == -1:
        return None
    
    frontmatter = '\n'.join(lines[1:end_idx])
    body = '\n'.join(lines[end_idx + 1:]).strip()
    
    data = _try_parse_yaml(frontmatter)
    if isinstance(data, dict) and body:
        # Add the markdown body as system_prompt if not already present
        if not data.get('system_prompt') and not data.get('prompt'):
            data['system_prompt'] = body
    
    return data


def _scan_static_data(root: Path) -> List[dict]:
    agents: List[dict] = []
    patterns = ["**/*.json", "**/*.yaml", "**/*.yml", "**/*.md"]
    for pattern in patterns:
        for fp in root.glob(pattern):
            if not fp.is_file():
                continue
            text = _read_text(fp)
            if not text:
                continue
            
            data = None
            if fp.suffix.lower() == '.md':
                data = _parse_markdown_frontmatter(text)
            else:
                data = _try_parse_json(text)
                if data is None:
                    data = _try_parse_yaml(text)
            
            if not isinstance(data, dict):
                continue
            # derive category from parent folder name as a fallback
            fallback_category = fp.parent.name
            agent = _extract_agent_from_obj(data, fallback_category, fp)
            if agent:
                agents.append(agent)
    return agents


def _maybe_snapshot_download_from_hf(url: str, target_dir: Path) -> Optional[Path]:
    # Attempt to fetch dataset to local static_data using huggingface_hub.
    try:
        from huggingface_hub import snapshot_download

        # Accept full URL like https://huggingface.co/datasets/owner/name
        m = re.match(r"https?://huggingface.co/datasets/([^/]+/[^/]+)", url.strip())
        if not m:
            return None
        repo_id = m.group(1)
        target_dir.mkdir(parents=True, exist_ok=True)
        local_path = snapshot_download(repo_id=repo_id, repo_type="dataset")
        # Mirror files into target_dir for predictable path
        src = Path(local_path)
        for p in src.rglob("*"):
            if p.is_file():
                rel = p.relative_to(src)
                dest = target_dir / rel
                dest.parent.mkdir(parents=True, exist_ok=True)
                try:
                    dest.write_bytes(p.read_bytes())
                except Exception:
                    pass
        return target_dir
    except Exception:
        # No network or hub not installed; just skip
        return None


def _parse_repo_id_from_url(url: str) -> Optional[str]:
    m = re.match(r"https?://huggingface.co/datasets/([^/]+/[^/]+)", url.strip())
    return m.group(1) if m else None


def _extract_agent_from_row(row: dict) -> Optional[dict]:
    if not isinstance(row, dict):
        return None
    name = _extract_field(row, ["name", "agent_name", "title", "id"]) or row.get("name")
    system_prompt = _extract_field(
        row,
        [
            "system_prompt",
            "prompt",
            "instructions",
            "system",
            "system_instructions",
            "system_text",
        ],
    )
    if not (name and system_prompt):
        return None
    description = _extract_field(row, ["description", "desc", "about", "summary"]) or ""
    category = _extract_field(row, ["category", "group", "type"]) or "uncategorized"
    category_slug = _slugify(category)
    agent_id = _slugify(f"{category_slug}-{name}")
    return {
        "id": agent_id,
        "name": name,
        "description": description,
        "system_prompt": system_prompt,
        "category": category_slug,
        "source": "hf-dataset-row",
    }


def _maybe_load_hf_dataset_rows(url: str) -> Optional[List[dict]]:
    try:
        import datasets  # type: ignore

        repo_id = _parse_repo_id_from_url(url)
        if not repo_id:
            return None

        # Try common splits; prefer train if present
        result: List[dict] = []
        loaded = datasets.load_dataset(repo_id)
        if isinstance(loaded, dict):
            split_order = ["train", "validation", "test"] + [k for k in loaded.keys() if k not in {"train", "validation", "test"}]
            for split in split_order:
                if split in loaded:
                    for row in loaded[split]:
                        a = _extract_agent_from_row(dict(row))
                        if a:
                            result.append(a)
        else:
            for row in loaded:  # type: ignore
                a = _extract_agent_from_row(dict(row))
                if a:
                    result.append(a)

        return result or None
    except Exception:
        return None


def load_agents() -> Tuple[Dict[str, Any], List[dict], List[str]]:
    """
    Returns (catalog_by_category, agents, warnings)
    catalog_by_category: {category_slug: {label: str, agents: [agent_id, ...]}}
    agents: list of agent dicts
    warnings: list of warning strings for UI
    """
    warnings: List[str] = []
    agents: List[dict] = []

    # Resolve datasource
    url = os.getenv("HF_DATASET_URL") or os.getenv("HF_DATASET_ID") or (_read_text(DATASOURCE_TXT) or "").strip()

    # 1) Prefer local static_data if present
    if STATIC_DATA_DIR.exists():
        agents = _scan_static_data(STATIC_DATA_DIR)
    # 2) Try to load dataset rows directly via datasets
    if not agents and url:
        maybe_agents = _maybe_load_hf_dataset_rows(url)
        if maybe_agents:
            agents = maybe_agents
    # 3) If rows failed, snapshot the repo and scan files
    if not agents and url:
        maybe_dir = _maybe_snapshot_download_from_hf(url, STATIC_DATA_DIR)
        if maybe_dir and maybe_dir.exists():
            agents = _scan_static_data(maybe_dir)
        else:
            warnings.append(
                "Dataset fetch unavailable. Add a local 'static_data' folder with agent configs."
            )
    if not url:
        warnings.append("No datasource URL found; using fallback sample data.")

    # 3) Fallback sample if nothing found
    if not agents:
        agents = [
            {
                "id": "code-assist-starter",
                "name": "Code Assist Starter",
                "description": "A simple code generation assistant for boilerplate tasks.",
                "system_prompt": (
                    "You are a helpful coding agent. Generate concise, correct code and explain key steps."
                ),
                "category": "code-assist",
                "source": "sample",
            },
            {
                "id": "docs-navigator",
                "name": "Docs Navigator",
                "description": "Answers questions using project docs and summarizes APIs.",
                "system_prompt": (
                    "Act as a technical writer. Read provided docs and produce accurate, concise answers with citations when possible."
                ),
                "category": "documentation",
                "source": "sample",
            },
        ]
        warnings.append(
            "Showing sample data. Add 'static_data' with JSON/YAML agent configs to replace."
        )

    # Dedupe by id, prefer first occurrence
    deduped: Dict[str, dict] = {}
    for a in agents:
        if isinstance(a, dict) and a.get("id") and a["id"] not in deduped:
            deduped[a["id"]] = a
    agents = list(deduped.values())

    # Build catalog
    catalog: Dict[str, Dict[str, Any]] = {}
    for a in agents:
        cat = a.get("category") or "uncategorized"
        cat_slug = _slugify(str(cat))
        if cat_slug not in catalog:
            catalog[cat_slug] = {"label": _titleize_slug(cat_slug), "agents": []}
        catalog[cat_slug]["agents"].append(a["id"])

    # Sort categories and agents by display name
    catalog = dict(sorted(catalog.items(), key=lambda kv: kv[1]["label"]))
    name_by_id = {a["id"]: a["name"] for a in agents}
    for c in catalog.values():
        c["agents"].sort(key=lambda aid: name_by_id.get(aid, aid).lower())

    return catalog, agents, warnings


# ----------------------
# Gradio application
# ----------------------


def build_ui():
    catalog, agents, warnings = load_agents()
    agent_by_id = {a["id"]: a for a in agents}

    # Initial selections
    first_cat = next(iter(catalog.keys())) if catalog else None
    first_agent = catalog[first_cat]["agents"][0] if first_cat and catalog[first_cat]["agents"] else None

    def show_about():
        """Return About page content."""
        about_content = """
        # About Code-Gen-Agents-Network
        
        This is a point-in-time network of code generation subagents created by **Daniel Rosehill**.
        
        The network represents a curated collection of specialized AI agents designed for various coding and development tasks. Each agent has been configured with specific system prompts and capabilities to assist with different aspects of software development.
        
        **Creator:** Daniel Rosehill  
        **Website:** [danielrosehill.com](https://danielrosehill.com)  
        **Dataset:** [Code-Gen-Agents-0925](https://huggingface.co/datasets/danielrosehill/Code-Gen-Agents-0925)
        
        ---
        
        ### Purpose
        
        This network serves as a comprehensive resource for developers looking to leverage specialized AI agents for:
        - Code generation and assistance
        - Documentation and writing tasks
        - Development workflow automation
        - Deployment and infrastructure management
        
        ### Usage
        
        Browse through the categories to find agents suited to your specific needs. Each agent includes:
        - A detailed description of its capabilities
        - The complete system prompt for implementation
        - Source information for reference
        
        Copy the system prompts to use these agents in your preferred AI interface or development environment.
        
        ### Data Source
        
        The agent configurations are sourced from the [Code-Gen-Agents-0925 dataset](https://huggingface.co/datasets/danielrosehill/Code-Gen-Agents-0925) on Hugging Face, which contains the complete collection of specialized coding agents and their system prompts.
        """
        return about_content

    def on_category_select(evt: gr.SelectData):
        """Handle category selection from the category list."""
        cat_slug = list(catalog.keys())[evt.index]
        agents_in_cat = catalog[cat_slug]["agents"]
        agent_choices = [agent_by_id[aid]["name"] for aid in agents_in_cat if aid in agent_by_id]
        first_agent_in_cat = agents_in_cat[0] if agents_in_cat else None
        
        # Update agent list and select first agent
        agent_update = gr.update(choices=agent_choices, value=agent_choices[0] if agent_choices else None)
        return agent_update, *on_agent_change(first_agent_in_cat)

    def on_agent_select(evt: gr.SelectData):
        """Handle agent selection from the agent list."""
        # Find which category is currently selected to get the right agent
        for cat_slug, cat_data in catalog.items():
            if evt.index < len(cat_data["agents"]):
                agent_id = cat_data["agents"][evt.index]
                return on_agent_change(agent_id)
        return on_agent_change(None)

    def on_agent_change(agent_id: Optional[str]):
        if not agent_id or agent_id not in agent_by_id:
            return (
                "# Select an agent",
                gr.update(value="", visible=True),
                gr.update(value="", visible=True),
                gr.update(value="", visible=False),
            )
        a = agent_by_id[agent_id]
        header = f"# {a['name']}"
        desc = a.get("description", "")
        prompt = a.get("system_prompt", "")
        src = a.get("source", "")
        footer = f"**Source:** {src}" if src else ""
        return (
            header,
            gr.update(value=desc, visible=True),
            gr.update(value=prompt, visible=True),
            gr.update(value=footer, visible=bool(footer)),
        )

    with gr.Blocks(title="Code Gen Agents Network", theme=gr.themes.Soft()) as demo:
        # Tab interface for About and Main content
        with gr.Tabs() as tabs:
            with gr.TabItem("Agents Network", id="main"):
                with gr.Row():
                    # Fixed-width left sidebar for categories
                    with gr.Column(scale=1, min_width=200):
                        gr.Markdown("### Categories")
                        category_list = gr.Radio(
                            choices=[catalog[slug]["label"] for slug in catalog.keys()],
                            value=catalog[first_cat]["label"] if first_cat else None,
                            label="",
                            interactive=True,
                            container=False
                        )
                        
                        if warnings:
                            with gr.Accordion("⚠️ Notes", open=False):
                                gr.Markdown("\n".join(f"- {w}" for w in warnings))

                    # Main content area
                    with gr.Column(scale=4):
                        # Top horizontal bar for agents
                        gr.Markdown("### Agents")
                        agent_list = gr.Radio(
                            choices=[agent_by_id[aid]["name"] for aid in catalog[first_cat]["agents"]] if first_cat else [],
                            value=agent_by_id[first_agent]["name"] if first_agent and first_agent in agent_by_id else None,
                            label="",
                            interactive=True,
                            container=False
                        )
                        
                        # Agent details below
                        md_header = gr.Markdown("# Select an agent")
                        
                        tb_desc = gr.Textbox(
                            label="Description",
                            lines=3,
                            show_copy_button=True,
                            interactive=False
                        )
                        
                        tb_prompt = gr.Textbox(
                            label="System Prompt",
                            lines=15,
                            show_copy_button=True,
                            interactive=False
                        )
                        md_footer = gr.Markdown(visible=False)
            
            with gr.TabItem("About", id="about"):
                about_markdown = gr.Markdown(show_about())

        # Wire events
        category_list.select(
            on_category_select,
            outputs=[agent_list, md_header, tb_desc, tb_prompt, md_footer]
        )
        
        # Handle agent selection - need to track current category
        current_category = gr.State(first_cat)
        
        def on_agent_radio_change(selected_agent_name, current_cat):
            if not selected_agent_name or not current_cat:
                return on_agent_change(None)
            
            # Find agent ID by name in current category
            for agent_id in catalog[current_cat]["agents"]:
                if agent_id in agent_by_id and agent_by_id[agent_id]["name"] == selected_agent_name:
                    return on_agent_change(agent_id)
            return on_agent_change(None)
        
        def update_current_category(evt: gr.SelectData):
            cat_slug = list(catalog.keys())[evt.index]
            agents_in_cat = catalog[cat_slug]["agents"]
            agent_choices = [agent_by_id[aid]["name"] for aid in agents_in_cat if aid in agent_by_id]
            first_agent_in_cat = agents_in_cat[0] if agents_in_cat else None
            
            return (
                cat_slug,
                gr.update(choices=agent_choices, value=agent_choices[0] if agent_choices else None),
                *on_agent_change(first_agent_in_cat)
            )
        
        category_list.select(
            update_current_category,
            outputs=[current_category, agent_list, md_header, tb_desc, tb_prompt, md_footer]
        )
        
        agent_list.change(
            on_agent_radio_change,
            inputs=[agent_list, current_category],
            outputs=[md_header, tb_desc, tb_prompt, md_footer]
        )

        # Initialize content
        if first_agent:
            header, desc_upd, prompt_upd, footer_upd = on_agent_change(first_agent)
            md_header.value = header
            tb_desc.value = desc_upd.value if hasattr(desc_upd, 'value') else ""
            tb_prompt.value = prompt_upd.value if hasattr(prompt_upd, 'value') else ""
            md_footer.value = footer_upd.value if hasattr(footer_upd, 'value') else ""
            md_footer.visible = bool(footer_upd.value) if hasattr(footer_upd, 'value') else False

    return demo


if __name__ == "__main__":
    demo = build_ui()
    demo.launch()