from DB import (
    Prompt,
    PromptCategory,
    Argument,
    User,
    get_session,
    ServerPrompt,
    ServerPromptCategory,
    ServerPromptArgument,
    CompanyPrompt,
    CompanyPromptCategory,
    CompanyPromptArgument,
    UserPromptOverride,
)
from Globals import DEFAULT_USER, get_default_user_id
from MagicalAuth import get_user_id, get_user_company_id
from SharedCache import shared_cache
from sqlalchemy.orm import joinedload
import os
import logging

THINK_ABOUT_IT = """## Available Context 
 {context} 

{COMMANDS}

Today's date is {date} .

Users latest input: {user_input} 

## System
The assistant's name is {agent_name}. The assistant is an expert at any given task due to its vast knowledge and wide array of available tools and in-context learning abilities.
Important elements of the assistant's personality and guidelines to follow are as follows:
The assistant thinks from first principles using available context to help enrich responses with learned information.
The assistant should be unapologetic about factual information and prioritize it over emotional sensitivity and social concerns in thought processes and responses.
The assistant does not mention that it is executing commands any more than a person would mention that they are breathing, this is just a passive ability that the assistant has and should use as it sees fit without asking permission, talking about what the commands do, or discussing which commands the assistant has access to. Command activities are logged passively without needing to be mentioned to the user.
The assistant is multimodal and has vision, hearing, speech, and image generation capabilities passively that also do not need to be mentioned to the user, these abilities manifest passively.
The assistant may use available context during the thought process to guide the response. Available context may contain the assistant's past memories, files sent from the user to the assistant, website content viewed by the assistant, conversation history between the user and the assistant, any other relevant information, or nothing at all.
The end goal will be to respond to the user's latest input using available context after going through the thought process.
Begin by enclosing all thoughts within <thinking> tags, exploring multiple angles and approaches through experimentation and reasoning.
The assistant can also theory craft ideas and explore different possibilities with code in <thinking> tags.
Break down the solution into clear steps within <step> tags. Use as many steps as needed — there is no step budget or limit. Complex problems may require many steps of deep thinking, experimentation, and iteration. Do not rush to a conclusion.
Use <step> tags to break down the solution into smaller, manageable steps, such as writing a single function if it is a coding task.
Continuously adjust your reasoning based on intermediate results and reflections, adapting your strategy as you progress.
Regularly evaluate progress using <reflection> tags. Be critical and honest about your reasoning process, consider where the assistant could be incorrect in thinking. If code was written in <thinking> or <step> tags, evaluate its correctness and efficiency within <reflection> tags.
Assign a quality score between 0.0 and 1.0 using <reward> tags after each reflection. Use this to guide your approach:

0.9+: Exceptional, continue with current approach
0.8-0.9: Minor refinements or additions can improve
0.5-0.8: Reflect and think deeper considering other approaches
Below 0.5: Seriously consider backtracking and trying a different approach

If unsure or if reward score is low, backtrack and try a different approach, explaining your decision within <thinking> tags.
For mathematical problems, show all work explicitly using LaTeX for formal notation and provide detailed proofs inside of a markdown code block.
Explore multiple solutions individually if possible, comparing approaches in reflections.
Use thoughts as a scratchpad, writing out all calculations and reasoning explicitly.
Conclude the thinking phase with a final reflection on the overall solution, discussing effectiveness, challenges, and solutions. Assign a final reward score, continue thinking and acting if the reward score is not 0.9+.
Synthesize the final answer within <answer> tags, providing a clear, concise but descriptive summary that captures the reasoned answer that does not leave out potentially important details.
The opening <answer> tag MUST be the first thing on a new line, not inline in a sentence and not inside a code block. This is the protocol boundary that tells AGiXT the final user-facing answer has started.
Thoughts are for the assistant and exist to give the assistant an opportunity to think through answers.
If any commands are available to the assistant for this task, they should only be executed before the <answer> tag. Another chain of logic will happen automatically when the assistant requests for a command to be executed.
Ensure that the <answer> tag is properly closed with </answer> at the end of the response that will go to the user.
Ensure that literal examples of <answer> tags inside sentences or markdown code blocks are never used as the real final answer boundary.
Never worry about running out of output tokens, you will be prompted to continue where you left off if you run out of output tokens. If you are prompted to continue, you can continue from where you left off without acknowledging the message or repeating anything that was already typed.
Thoughts and reflections will not be available to the user, only to the assistant from memories in future interactions where relevant.
The content of the <answer> block should be in markdown.
In the answer block, never ask the user to hold while working on something - all work should happen during the thinking/reflecting/executing phase, the answer block should only contain the final output to the user.
**Only the <answer> tag will be available to the user, so make sure it is complete, well-reasoned, accurate and ends in </answer>. It must end in </answer> !**

## Coding & Code Repository Tasks
If the available context indicates this task involves reading or modifying source code in a repository (signals: file/terminal/git/`gh` commands are available, the workspace contains source files, the user mentions a repo/branch/PR/file/function/test/build), shift into the deliberate coding discipline below. For non-coding tasks, ignore this section.

The single rule: **be curious before being productive.** Most low-quality coding output comes from editing the first plausible file without understanding how it fits in. Move through these phases in order; loop back when a later phase invalidates an earlier assumption, but never silently drift.

1. **Maintain `context.md` in the workspace root as the source of truth.** Read it first if it exists; create it if it does not. It holds, in this order: `Task`, `Repository map`, `Relevant files & symbols`, `Constraints & conventions`, `Open questions`, `Plan`, `Discoveries during build`, `Verification`. Update it after every meaningful discovery — append new findings, never silently overwrite history.

2. **Explore.** Before any production edit, use `List Directory`, `Glob File Search`, `Grep Search`, `Search File Content`, `Find Symbol Usages`, `Find References`, `Read File`, and `Create or Update Codebase Map` to read the files that will change, the files that call them, and the tests that cover them. Identify the project's build, lint, and test commands from `package.json`, `pyproject.toml`, `Makefile`, `AGENTS.md`, `README.md`, `.github/workflows/`, etc. Populate `context.md` through `Constraints & conventions`.

3. **Plan.** Write a numbered plan into `context.md`. Each step names specific files, functions, and tests, with an obvious success criterion. Mirror the plan into the conversation's todo list with `Create Todo Items in Bulk` so progress is tracked in the database, not just in markdown.

4. **Build.** Execute the plan step by step, marking todos in-progress and completed as you go. Prefer `Modify File`, `Insert in File`, `Multi-File Replace`, and `Search and Replace Regex` for surgical edits; use `Write to File` only for new files. After each step, briefly verify (re-read the region, run a targeted test, syntax check) before moving on. Discoveries that change the plan get appended to `Discoveries during build` in `context.md`, and the `Plan` section is updated.

5. **Reflect & test.** Before declaring done, deliberately ask: what did I assume during planning that I now know to be wrong? What edge cases did I skip? Append findings to `context.md`. Then actually run the project's tests and linter/formatter via `Use Terminal in Workspace` (e.g. `pytest`, `npm test`, `cargo test`, `black .`, `ruff check .`, `eslint .`, `prettier --check .`). Record every command and its result under `Verification`. If anything fails, fix it and re-run until clean.

6. **Branch & pull request — never commit to `main` or `master`.** Create or reuse a descriptive feature branch: `git checkout -b <type>/<short-slug>` where `<type>` is `feat`, `fix`, `chore`, `refactor`, `docs`, or `test`. Commit with a clear conventional message. Push with `git push -u origin <branch>`. If no PR exists for this branch, open one with `gh pr create` against the repository's default branch (look it up with `gh repo view --json defaultBranchRef -q .defaultBranchRef.name`); the PR body must contain `Summary`, `Plan executed` (with checkmarks), `Verification` (commands + results), and `Notes / follow-ups`, all derived from `context.md`. If a PR already exists for the branch, push the new commits and post a `gh pr comment` summarizing the latest changes — do NOT open a duplicate PR.

The `<answer>` for a coding task should include: a plain-language summary of what was done, the branch name and PR URL (if applicable), any deferred follow-ups, and the path to `context.md` so the user can audit the reasoning trail. The full `Coding Workflow` chain enforces these phases as discrete steps and can be invoked when the task warrants it.
"""


class Prompts:
    def __init__(self, user=DEFAULT_USER):
        self.user = user
        self.user_id = get_user_id(user)
        self.company_id = get_user_company_id(user)

    def _get_prompt_args_from_content(self, prompt_text):
        """Extract {variable} arguments from prompt content."""
        prompt_args = []
        start_index = prompt_text.find("{")
        while start_index != -1:
            end_index = prompt_text.find("}", start_index)
            if end_index != -1:
                prompt_args.append(prompt_text[start_index + 1 : end_index])
                start_index = prompt_text.find("{", end_index)
            else:
                break
        return prompt_args

    # Alias for backwards compatibility
    def get_prompt_args(self, prompt_text):
        return self._get_prompt_args_from_content(prompt_text)

    # =========================================================================
    # User-level prompt methods (existing behavior with tiered fallback)
    # =========================================================================

    def add_prompt(self, prompt_name, prompt, prompt_category="Default"):
        """Add a user-level prompt."""
        session = get_session()
        if not prompt_category or prompt_category == "":
            prompt_category = "Default"
        pc = (
            session.query(PromptCategory)
            .filter(
                PromptCategory.name == prompt_category,
                PromptCategory.user_id == self.user_id,
            )
            .first()
        )
        if not pc:
            pc = PromptCategory(
                name=prompt_category,
                description=f"{prompt_category} category",
                user_id=self.user_id,
            )
            session.add(pc)
            session.commit()
        prompt_obj = Prompt(
            name=prompt_name,
            description="",
            content=prompt,
            prompt_category_id=pc.id,
            user_id=self.user_id,
        )
        session.add(prompt_obj)
        session.commit()

        # Populate prompt arguments
        prompt_args = self._get_prompt_args_from_content(prompt)
        for arg in prompt_args:
            argument = Argument(
                prompt_id=prompt_obj.id,
                name=arg,
            )
            session.add(argument)
        session.commit()
        prompt_id = str(prompt_obj.id)
        session.close()
        return prompt_id

    def get_prompt(self, prompt_name: str, prompt_category: str = "Default"):
        """
        Get a prompt by name with tiered resolution:
        1. User-level prompt (highest priority)
        2. Company-level prompt
        3. Server-level prompt (global)
        4. Legacy file-based prompt (lowest priority, imports to user level)

        For internal prompts, server level is always used.
        Results are cached for 60s to avoid repeated DB queries on hot paths.
        """
        if prompt_name.lower() == "think about it":
            return THINK_ABOUT_IT
        cache_key = (
            f"prompt:{self.user_id}:{self.company_id}:{prompt_category}:{prompt_name}"
        )
        cached = shared_cache.get(cache_key)
        if cached is not None:
            return cached

        session = get_session()

        # 1. Try user-level prompt first
        prompt = (
            session.query(Prompt)
            .join(PromptCategory)
            .filter(
                Prompt.name == prompt_name,
                Prompt.user_id == self.user_id,
                PromptCategory.name == prompt_category,
            )
            .first()
        )
        if prompt:
            content = prompt.content
            session.close()
            shared_cache.set(cache_key, content, 60)
            return content

        # 2. Try company-level prompt
        if self.company_id:
            company_prompt = (
                session.query(CompanyPrompt)
                .join(CompanyPromptCategory)
                .filter(
                    CompanyPrompt.name == prompt_name,
                    CompanyPrompt.company_id == self.company_id,
                    CompanyPromptCategory.name == prompt_category,
                )
                .first()
            )
            if company_prompt:
                content = company_prompt.content
                session.close()
                shared_cache.set(cache_key, content, 60)
                return content

        # 3. Try server-level prompt (including internal prompts)
        server_prompt = (
            session.query(ServerPrompt)
            .join(ServerPromptCategory)
            .filter(
                ServerPrompt.name == prompt_name,
                ServerPromptCategory.name == prompt_category,
            )
            .first()
        )
        if server_prompt:
            content = server_prompt.content
            session.close()
            shared_cache.set(cache_key, content, 60)
            return content

        # 4. Legacy: Try default user's prompts (for backwards compatibility)
        default_uid = get_default_user_id()
        if default_uid:
            default_prompt = (
                session.query(Prompt)
                .join(PromptCategory)
                .filter(
                    Prompt.name == prompt_name,
                    Prompt.user_id == default_uid,
                    PromptCategory.name == prompt_category,
                )
                .first()
            )
            if default_prompt:
                content = default_prompt.content
                session.close()
                shared_cache.set(cache_key, content, 60)
                return content

        # 5. Legacy: Try file-based prompts and import to user level
        prompt_file = os.path.normpath(
            os.path.join(os.getcwd(), "prompts", prompt_category, f"{prompt_name}.txt")
        )
        base_path = os.path.join(os.getcwd(), "prompts")
        if prompt_file.startswith(base_path) and os.path.exists(prompt_file):
            with open(prompt_file, "r") as f:
                prompt_content = f.read()
            session.close()
            # Import to user level
            self.add_prompt(
                prompt_name=prompt_name,
                prompt=prompt_content,
                prompt_category=prompt_category,
            )
            shared_cache.set(cache_key, prompt_content, 60)
            return prompt_content

        session.close()
        return None

    def get_user_prompts(self):
        """
        Get all prompts available to the user from all tiers.
        Returns prompts with source indicators (server/company/user).
        User prompts override company, which override server prompts of the same name.
        """
        session = get_session()
        prompts_dict = {}  # name -> prompt data (for deduplication)

        # 1. Server-level prompts (non-internal only)
        server_prompts = (
            session.query(ServerPrompt)
            .options(
                joinedload(ServerPrompt.arguments),
                joinedload(ServerPrompt.category),
            )
            .join(ServerPromptCategory)
            .filter(ServerPrompt.is_internal == False)
            .all()
        )
        for prompt in server_prompts:
            try:
                prompt_args = [arg.name for arg in prompt.arguments]
            except:
                prompt_args = []
            prompts_dict[prompt.name] = {
                "id": str(prompt.id),
                "name": prompt.name,
                "category": prompt.category.name,
                "content": prompt.content,
                "description": prompt.description,
                "arguments": prompt_args,
                "source": "server",
                "is_override": False,
            }

        # 2. Company-level prompts (override server)
        if self.company_id:
            company_prompts = (
                session.query(CompanyPrompt)
                .options(
                    joinedload(CompanyPrompt.arguments),
                    joinedload(CompanyPrompt.category),
                )
                .filter(CompanyPrompt.company_id == self.company_id)
                .all()
            )
            for prompt in company_prompts:
                try:
                    prompt_args = [arg.name for arg in prompt.arguments]
                except:
                    prompt_args = []
                is_override = prompt.server_prompt_id is not None
                prompts_dict[prompt.name] = {
                    "id": str(prompt.id),
                    "name": prompt.name,
                    "category": prompt.category.name,
                    "content": prompt.content,
                    "description": prompt.description,
                    "arguments": prompt_args,
                    "source": "company",
                    "is_override": is_override,
                    "parent_id": (
                        str(prompt.server_prompt_id)
                        if prompt.server_prompt_id
                        else None
                    ),
                }

        # 3. User-level prompts (override company and server)
        user_prompts = (
            session.query(Prompt)
            .options(
                joinedload(Prompt.arguments),
                joinedload(Prompt.prompt_category),
            )
            .filter(Prompt.user_id == self.user_id)
            .all()
        )
        # Batch-fetch all user prompt overrides in a single query (N+1 → 1)
        all_user_overrides = (
            session.query(UserPromptOverride)
            .filter(UserPromptOverride.user_id == self.user_id)
            .all()
        )
        override_by_prompt_id = {str(o.prompt_id): o for o in all_user_overrides}

        for prompt in user_prompts:
            try:
                prompt_args = [arg.name for arg in prompt.arguments]
            except:
                prompt_args = []

            # Check if this is an override (from pre-fetched dict)
            override = override_by_prompt_id.get(str(prompt.id))
            is_override = override is not None
            parent_id = None
            if override:
                if override.server_prompt_id:
                    parent_id = str(override.server_prompt_id)
                elif override.company_prompt_id:
                    parent_id = str(override.company_prompt_id)

            prompts_dict[prompt.name] = {
                "id": str(prompt.id),
                "name": prompt.name,
                "category": prompt.prompt_category.name,
                "content": prompt.content,
                "description": prompt.description,
                "arguments": prompt_args,
                "source": "user",
                "is_override": is_override,
                "parent_id": parent_id,
            }

        # 4. Legacy: Include global prompts from DEFAULT_USER that user doesn't have
        default_uid = get_default_user_id()
        if default_uid and default_uid != str(self.user_id):
            global_prompts = (
                session.query(Prompt)
                .options(
                    joinedload(Prompt.arguments),
                    joinedload(Prompt.prompt_category),
                )
                .filter(Prompt.user_id == default_uid)
                .all()
            )
            for prompt in global_prompts:
                if prompt.name not in prompts_dict:
                    try:
                        prompt_args = [arg.name for arg in prompt.arguments]
                    except:
                        prompt_args = []
                    prompts_dict[prompt.name] = {
                        "id": str(prompt.id),
                        "name": prompt.name,
                        "category": prompt.prompt_category.name,
                        "content": prompt.content,
                        "description": prompt.description,
                        "arguments": prompt_args,
                        "source": "global",  # Legacy global
                        "is_override": False,
                    }

        session.close()
        return list(prompts_dict.values())

    def get_global_prompts(self):
        """
        Get server-level and legacy global prompts that the user doesn't have overridden.
        """
        session = get_session()
        prompts = []

        # Get user's prompt names for exclusion
        user_prompt_names = {
            row[0]
            for row in session.query(Prompt.name)
            .filter(Prompt.user_id == self.user_id)
            .all()
        }

        # Server-level prompts (non-internal)
        server_prompts = (
            session.query(ServerPrompt)
            .options(
                joinedload(ServerPrompt.arguments),
                joinedload(ServerPrompt.category),
            )
            .filter(ServerPrompt.is_internal == False)
            .all()
        )
        for prompt in server_prompts:
            if prompt.name not in user_prompt_names:
                try:
                    prompt_args = [arg.name for arg in prompt.arguments]
                except:
                    prompt_args = []
                prompts.append(
                    {
                        "id": str(prompt.id),
                        "name": prompt.name,
                        "category": prompt.category.name,
                        "content": prompt.content,
                        "description": prompt.description,
                        "arguments": prompt_args,
                        "source": "server",
                    }
                )

        # Company-level prompts
        if self.company_id:
            company_prompts = (
                session.query(CompanyPrompt)
                .options(
                    joinedload(CompanyPrompt.arguments),
                    joinedload(CompanyPrompt.category),
                )
                .filter(CompanyPrompt.company_id == self.company_id)
                .all()
            )
            for prompt in company_prompts:
                if prompt.name not in user_prompt_names:
                    try:
                        prompt_args = [arg.name for arg in prompt.arguments]
                    except:
                        prompt_args = []
                    prompts.append(
                        {
                            "id": str(prompt.id),
                            "name": prompt.name,
                            "category": prompt.category.name,
                            "content": prompt.content,
                            "description": prompt.description,
                            "arguments": prompt_args,
                            "source": "company",
                        }
                    )

        # Legacy global prompts
        default_uid = get_default_user_id()
        if default_uid:
            global_prompts = (
                session.query(Prompt)
                .options(
                    joinedload(Prompt.arguments),
                    joinedload(Prompt.prompt_category),
                )
                .filter(Prompt.user_id == default_uid)
                .all()
            )
            for prompt in global_prompts:
                if prompt.name not in user_prompt_names:
                    try:
                        prompt_args = [arg.name for arg in prompt.arguments]
                    except:
                        prompt_args = []
                    prompts.append(
                        {
                            "id": str(prompt.id),
                            "name": prompt.name,
                            "category": prompt.prompt_category.name,
                            "content": prompt.content,
                            "description": prompt.description,
                            "arguments": prompt_args,
                            "source": "global",
                        }
                    )

        session.close()
        return prompts

    def get_prompts_markdown(self):
        """Get all prompts as markdown documentation."""
        prompt_data = self.get_user_prompts()
        prompt_markdown = ""
        for prompt in prompt_data:
            prompt_markdown += f"## Prompt Name: `{prompt['name']}`\n"
            if prompt.get("description"):
                prompt_markdown += f"**Description:** {prompt['description']}\n"
            prompt_markdown += (
                f"**Arguments:** {', '.join(prompt.get('arguments', []))}\n"
            )
            prompt_markdown += f"**Source:** {prompt.get('source', 'user')}\n\n"
        return prompt_markdown

    def get_prompts(self, prompt_category="Default", include_ids: bool = False):
        """Get prompts in a category with tiered resolution."""
        if not prompt_category:
            prompt_category = "Default"
        session = get_session()
        prompts_dict = {}

        # 1. Server-level prompts in category (non-internal)
        server_prompts = (
            session.query(ServerPrompt)
            .join(ServerPromptCategory)
            .filter(
                ServerPromptCategory.name == prompt_category,
                ServerPrompt.is_internal == False,
            )
            .all()
        )
        for prompt in server_prompts:
            if include_ids:
                prompts_dict[prompt.name] = {
                    "id": str(prompt.id),
                    "name": prompt.name,
                    "source": "server",
                }
            else:
                prompts_dict[prompt.name] = prompt.name

        # 2. Company-level prompts in category
        if self.company_id:
            company_prompts = (
                session.query(CompanyPrompt)
                .join(CompanyPromptCategory)
                .filter(
                    CompanyPromptCategory.name == prompt_category,
                    CompanyPrompt.company_id == self.company_id,
                )
                .all()
            )
            for prompt in company_prompts:
                if include_ids:
                    prompts_dict[prompt.name] = {
                        "id": str(prompt.id),
                        "name": prompt.name,
                        "source": "company",
                    }
                else:
                    prompts_dict[prompt.name] = prompt.name

        # 3. User-level prompts in category
        user_prompts = (
            session.query(Prompt)
            .join(PromptCategory)
            .filter(
                PromptCategory.name == prompt_category,
                Prompt.user_id == self.user_id,
            )
            .all()
        )
        for prompt in user_prompts:
            if include_ids:
                prompts_dict[prompt.name] = {
                    "id": str(prompt.id),
                    "name": prompt.name,
                    "source": "user",
                }
            else:
                prompts_dict[prompt.name] = prompt.name

        # 4. Legacy global prompts
        default_uid = get_default_user_id()
        if default_uid:
            global_prompts = (
                session.query(Prompt)
                .join(PromptCategory)
                .filter(
                    PromptCategory.name == prompt_category,
                    Prompt.user_id == default_uid,
                )
                .all()
            )
            for prompt in global_prompts:
                if prompt.name not in prompts_dict:
                    if include_ids:
                        prompts_dict[prompt.name] = {
                            "id": str(prompt.id),
                            "name": prompt.name,
                            "source": "global",
                        }
                    else:
                        prompts_dict[prompt.name] = prompt.name

        session.close()
        return list(prompts_dict.values())

    def get_prompts_by_category_id(self, category_id: str):
        """Get prompts by category ID with full details including ID."""
        session = get_session()
        prompts = []

        # Check if this is a server category
        server_category = (
            session.query(ServerPromptCategory)
            .filter(ServerPromptCategory.id == category_id)
            .first()
        )
        if server_category:
            for prompt in server_category.prompts:
                if not prompt.is_internal:
                    prompts.append(
                        {"id": str(prompt.id), "name": prompt.name, "source": "server"}
                    )
            session.close()
            return prompts

        # Check if this is a company category
        company_category = (
            session.query(CompanyPromptCategory)
            .filter(
                CompanyPromptCategory.id == category_id,
                CompanyPromptCategory.company_id == self.company_id,
            )
            .first()
        )
        if company_category:
            for prompt in company_category.prompts:
                prompts.append(
                    {"id": str(prompt.id), "name": prompt.name, "source": "company"}
                )
            session.close()
            return prompts

        # Must be a user category
        default_uid = get_default_user_id()
        global_prompts = (
            (
                session.query(Prompt)
                .filter(
                    Prompt.user_id == default_uid,
                    Prompt.prompt_category_id == category_id,
                )
                .all()
            )
            if default_uid
            else []
        )
        user_prompts = (
            session.query(Prompt)
            .filter(
                Prompt.user_id == self.user_id,
                Prompt.prompt_category_id == category_id,
            )
            .all()
        )

        seen_names = {}
        for prompt in global_prompts:
            seen_names[prompt.name] = {
                "id": str(prompt.id),
                "name": prompt.name,
                "source": "global",
            }
        for prompt in user_prompts:
            seen_names[prompt.name] = {
                "id": str(prompt.id),
                "name": prompt.name,
                "source": "user",
            }

        session.close()
        return list(seen_names.values())

    def delete_prompt(self, prompt_name, prompt_category="Default"):
        """Delete a user-level prompt."""
        if not prompt_category:
            prompt_category = "Default"
        session = get_session()
        prompt = (
            session.query(Prompt)
            .filter_by(name=prompt_name)
            .join(PromptCategory)
            .filter(
                PromptCategory.name == prompt_category, Prompt.user_id == self.user_id
            )
            .first()
        )
        if prompt:
            # Also delete any override tracking
            session.query(UserPromptOverride).filter(
                UserPromptOverride.prompt_id == prompt.id
            ).delete()
            session.delete(prompt)
            session.commit()
        session.close()
        # Invalidate cached prompt content
        shared_cache.delete(
            f"prompt:{self.user_id}:{self.company_id}:{prompt_category}:{prompt_name}"
        )

    def update_prompt(self, prompt_name, prompt, prompt_category="Default"):
        """Update a user-level prompt or clone from parent tier if editing inherited."""
        if not prompt_category:
            prompt_category = "Default"
        session = get_session()

        # Check if user has this prompt
        prompt_obj = (
            session.query(Prompt)
            .filter(
                Prompt.name == prompt_name,
                Prompt.user_id == self.user_id,
                Prompt.prompt_category.has(name=prompt_category),
            )
            .first()
        )

        if prompt_obj:
            # User owns this prompt, update it directly
            prompt_obj.content = prompt
            session.commit()
            # Update arguments
            prompt_args = self._get_prompt_args_from_content(prompt)
            existing_args = (
                session.query(Argument).filter_by(prompt_id=prompt_obj.id).all()
            )
            existing_arg_names = {arg.name for arg in existing_args}
            for arg in existing_args:
                if arg.name not in prompt_args:
                    session.delete(arg)
            for arg in prompt_args:
                if arg not in existing_arg_names:
                    session.add(Argument(prompt_id=prompt_obj.id, name=arg))
            session.commit()
            session.close()
            # Invalidate cached prompt content
            shared_cache.delete(
                f"prompt:{self.user_id}:{self.company_id}:{prompt_category}:{prompt_name}"
            )
            return str(prompt_obj.id)

        # User doesn't have this prompt - need to clone from parent tier
        session.close()
        # Invalidate cached prompt content
        shared_cache.delete(
            f"prompt:{self.user_id}:{self.company_id}:{prompt_category}:{prompt_name}"
        )
        return self._clone_and_edit_prompt(prompt_name, prompt, prompt_category)

    def _clone_and_edit_prompt(self, prompt_name, new_content, prompt_category):
        """Clone a prompt from parent tier (server/company) to user level for editing."""
        session = get_session()

        server_prompt_id = None
        company_prompt_id = None
        source_prompt = None

        # Check company-level first
        if self.company_id:
            company_prompt = (
                session.query(CompanyPrompt)
                .join(CompanyPromptCategory)
                .filter(
                    CompanyPrompt.name == prompt_name,
                    CompanyPrompt.company_id == self.company_id,
                    CompanyPromptCategory.name == prompt_category,
                )
                .first()
            )
            if company_prompt:
                source_prompt = company_prompt
                company_prompt_id = company_prompt.id

        # Check server-level
        if not source_prompt:
            server_prompt = (
                session.query(ServerPrompt)
                .join(ServerPromptCategory)
                .filter(
                    ServerPrompt.name == prompt_name,
                    ServerPromptCategory.name == prompt_category,
                )
                .first()
            )
            if server_prompt:
                source_prompt = server_prompt
                server_prompt_id = server_prompt.id

        if not source_prompt:
            session.close()
            # No parent prompt found, create new
            return self.add_prompt(prompt_name, new_content, prompt_category)

        # Create user's copy
        session.close()
        new_prompt_id = self.add_prompt(prompt_name, new_content, prompt_category)

        # Track the override
        session = get_session()
        override = UserPromptOverride(
            user_id=self.user_id,
            prompt_id=new_prompt_id,
            server_prompt_id=server_prompt_id,
            company_prompt_id=company_prompt_id,
        )
        session.add(override)
        session.commit()
        session.close()

        return new_prompt_id

    def revert_prompt_to_default(self, prompt_id: str):
        """
        Revert a user's customized prompt back to the parent (server/company) version.
        Deletes the user's override and removes the tracking record.
        """
        session = get_session()

        override = (
            session.query(UserPromptOverride)
            .filter(
                UserPromptOverride.user_id == self.user_id,
                UserPromptOverride.prompt_id == prompt_id,
            )
            .first()
        )

        if not override:
            session.close()
            return {"success": False, "message": "This prompt is not an override"}

        # Delete the user's prompt
        prompt = session.query(Prompt).filter(Prompt.id == prompt_id).first()
        if prompt:
            session.delete(prompt)

        # Delete the override tracking
        session.delete(override)
        session.commit()
        session.close()

        return {"success": True, "message": "Reverted to default"}

    def rename_prompt(self, prompt_name, new_prompt_name, prompt_category="Default"):
        """Rename a user-level prompt."""
        if not prompt_category:
            prompt_category = "Default"
        session = get_session()
        prompt = (
            session.query(Prompt)
            .filter(
                Prompt.name == prompt_name,
                Prompt.user_id == self.user_id,
                Prompt.prompt_category.has(name=prompt_category),
            )
            .join(PromptCategory)
            .filter(
                PromptCategory.name == prompt_category, Prompt.user_id == self.user_id
            )
            .first()
        )
        if prompt:
            prompt.name = new_prompt_name
            session.commit()
        session.close()

    def get_prompt_categories(self, include_ids: bool = False):
        """Get all prompt categories from all tiers."""
        session = get_session()
        categories_dict = {}

        # 1. Server categories (non-internal)
        server_categories = (
            session.query(ServerPromptCategory)
            .filter(ServerPromptCategory.is_internal == False)
            .all()
        )
        for cat in server_categories:
            if include_ids:
                categories_dict[cat.name] = {
                    "id": str(cat.id),
                    "name": cat.name,
                    "description": cat.description,
                    "source": "server",
                }
            else:
                categories_dict[cat.name] = cat.name

        # 2. Company categories
        if self.company_id:
            company_categories = (
                session.query(CompanyPromptCategory)
                .filter(CompanyPromptCategory.company_id == self.company_id)
                .all()
            )
            for cat in company_categories:
                if include_ids:
                    categories_dict[cat.name] = {
                        "id": str(cat.id),
                        "name": cat.name,
                        "description": cat.description,
                        "source": "company",
                    }
                else:
                    categories_dict[cat.name] = cat.name

        # 3. User categories
        user_categories = (
            session.query(PromptCategory)
            .filter(PromptCategory.user_id == self.user_id)
            .all()
        )
        for cat in user_categories:
            if include_ids:
                categories_dict[cat.name] = {
                    "id": str(cat.id),
                    "name": cat.name,
                    "description": cat.description,
                    "source": "user",
                }
            else:
                categories_dict[cat.name] = cat.name

        # 4. Legacy global categories
        default_uid = get_default_user_id()
        if default_uid:
            global_categories = (
                session.query(PromptCategory)
                .filter(PromptCategory.user_id == default_uid)
                .all()
            )
            for cat in global_categories:
                if cat.name not in categories_dict:
                    if include_ids:
                        categories_dict[cat.name] = {
                            "id": str(cat.id),
                            "name": cat.name,
                            "description": cat.description,
                            "source": "global",
                        }
                    else:
                        categories_dict[cat.name] = cat.name

        session.close()
        return list(categories_dict.values())

    def get_prompt_by_id(self, prompt_id: str):
        """Get a prompt's content by ID across all tiers."""
        session = get_session()

        # Try user prompt
        prompt = (
            session.query(Prompt)
            .filter(
                Prompt.id == prompt_id,
                Prompt.user_id == self.user_id,
            )
            .first()
        )
        if prompt:
            session.close()
            if prompt.name.lower() == "think about it":
                return THINK_ABOUT_IT
            return prompt.content

        # Try company prompt
        if self.company_id:
            company_prompt = (
                session.query(CompanyPrompt)
                .filter(
                    CompanyPrompt.id == prompt_id,
                    CompanyPrompt.company_id == self.company_id,
                )
                .first()
            )
            if company_prompt:
                session.close()
                return company_prompt.content

        # Try server prompt
        server_prompt = (
            session.query(ServerPrompt).filter(ServerPrompt.id == prompt_id).first()
        )
        if server_prompt:
            session.close()
            return server_prompt.content

        # Try legacy global prompt
        default_uid = get_default_user_id()
        if default_uid:
            global_prompt = (
                session.query(Prompt)
                .filter(
                    Prompt.id == prompt_id,
                    Prompt.user_id == default_uid,
                )
                .first()
            )
            if global_prompt:
                session.close()
                return global_prompt.content

        session.close()
        return None

    def update_prompt_by_id(
        self, prompt_id: str, prompt_name: str = None, prompt: str = None
    ):
        """Update a prompt by ID. If it's a parent prompt, clone it first."""
        session = get_session()

        # Check if user owns this prompt
        prompt_obj = (
            session.query(Prompt)
            .filter(
                Prompt.id == prompt_id,
                Prompt.user_id == self.user_id,
            )
            .first()
        )

        if prompt_obj:
            # User owns it, update directly
            if prompt_name is not None:
                prompt_obj.name = prompt_name
            if prompt is not None:
                prompt_obj.content = prompt
                # Update arguments
                session.query(Argument).filter(Argument.prompt_id == prompt_id).delete()
                for arg in self._get_prompt_args_from_content(prompt):
                    session.add(Argument(prompt_id=prompt_id, name=arg))
            updated_prompt_id = str(prompt_obj.id)
            session.commit()
            session.close()
            # Invalidate all cached prompts for this user
            shared_cache.delete_pattern(f"prompt:{self.user_id}:*")
            return updated_prompt_id

        # Get the source prompt for cloning
        source_content = None
        source_name = None
        source_category = "Default"
        server_prompt_id = None
        company_prompt_id = None

        # Check company level
        if self.company_id:
            company_prompt = (
                session.query(CompanyPrompt)
                .filter(
                    CompanyPrompt.id == prompt_id,
                    CompanyPrompt.company_id == self.company_id,
                )
                .first()
            )
            if company_prompt:
                source_content = prompt if prompt else company_prompt.content
                source_name = prompt_name if prompt_name else company_prompt.name
                source_category = company_prompt.category.name
                company_prompt_id = company_prompt.id

        # Check server level
        if not source_content:
            server_prompt = (
                session.query(ServerPrompt).filter(ServerPrompt.id == prompt_id).first()
            )
            if server_prompt:
                source_content = prompt if prompt else server_prompt.content
                source_name = prompt_name if prompt_name else server_prompt.name
                source_category = server_prompt.category.name
                server_prompt_id = server_prompt.id

        session.close()

        if not source_content:
            raise Exception("Prompt not found")

        # Create user copy
        new_prompt_id = self.add_prompt(source_name, source_content, source_category)

        # Track override
        session = get_session()
        override = UserPromptOverride(
            user_id=self.user_id,
            prompt_id=new_prompt_id,
            server_prompt_id=server_prompt_id,
            company_prompt_id=company_prompt_id,
        )
        session.add(override)
        session.commit()
        session.close()

        # Invalidate all cached prompts for this user
        shared_cache.delete_pattern(f"prompt:{self.user_id}:*")

        return new_prompt_id

    def delete_prompt_by_id(self, prompt_id: str):
        """Delete a user's prompt by ID."""
        session = get_session()
        prompt = (
            session.query(Prompt)
            .filter(
                Prompt.id == prompt_id,
                Prompt.user_id == self.user_id,
            )
            .first()
        )
        if not prompt:
            session.close()
            raise Exception(
                "Prompt not found or you don't have permission to delete it"
            )

        # Delete override tracking if exists
        session.query(UserPromptOverride).filter(
            UserPromptOverride.prompt_id == prompt_id
        ).delete()

        session.delete(prompt)
        session.commit()
        session.close()
        # Invalidate all cached prompts for this user
        shared_cache.delete_pattern(f"prompt:{self.user_id}:*")

    def get_prompt_details_by_id(self, prompt_id: str):
        """Get full prompt details by ID across all tiers."""
        session = get_session()

        # Try user prompt
        prompt = (
            session.query(Prompt)
            .filter(
                Prompt.id == prompt_id,
                Prompt.user_id == self.user_id,
            )
            .first()
        )
        if prompt:
            try:
                prompt_args = [arg.name for arg in prompt.arguments]
            except:
                prompt_args = []

            # Check if this is an override
            override = (
                session.query(UserPromptOverride)
                .filter(UserPromptOverride.prompt_id == prompt_id)
                .first()
            )

            result = {
                "prompt_name": prompt.name,
                "prompt_category": prompt.prompt_category.name,
                "prompt": prompt.content,
                "description": prompt.description,
                "arguments": prompt_args,
                "id": str(prompt.id),
                "source": "user",
                "is_override": override is not None,
                "can_revert": override is not None,
            }
            session.close()
            return result

        # Try company prompt
        if self.company_id:
            company_prompt = (
                session.query(CompanyPrompt)
                .filter(
                    CompanyPrompt.id == prompt_id,
                    CompanyPrompt.company_id == self.company_id,
                )
                .first()
            )
            if company_prompt:
                try:
                    prompt_args = [arg.name for arg in company_prompt.arguments]
                except:
                    prompt_args = []
                result = {
                    "prompt_name": company_prompt.name,
                    "prompt_category": company_prompt.category.name,
                    "prompt": company_prompt.content,
                    "description": company_prompt.description,
                    "arguments": prompt_args,
                    "id": str(company_prompt.id),
                    "source": "company",
                    "is_override": company_prompt.server_prompt_id is not None,
                    "can_revert": False,  # Would need company admin to revert
                }
                session.close()
                return result

        # Try server prompt
        server_prompt = (
            session.query(ServerPrompt).filter(ServerPrompt.id == prompt_id).first()
        )
        if server_prompt:
            try:
                prompt_args = [arg.name for arg in server_prompt.arguments]
            except:
                prompt_args = []
            result = {
                "prompt_name": server_prompt.name,
                "prompt_category": server_prompt.category.name,
                "prompt": server_prompt.content,
                "description": server_prompt.description,
                "arguments": prompt_args,
                "id": str(server_prompt.id),
                "source": "server",
                "is_override": False,
                "can_revert": False,
            }
            session.close()
            return result

        # Try legacy global
        user_data = session.query(User).filter(User.email == DEFAULT_USER).first()
        if user_data:
            global_prompt = (
                session.query(Prompt)
                .filter(
                    Prompt.id == prompt_id,
                    Prompt.user_id == user_data.id,
                )
                .first()
            )
            if global_prompt:
                try:
                    prompt_args = [arg.name for arg in global_prompt.arguments]
                except:
                    prompt_args = []
                result = {
                    "prompt_name": global_prompt.name,
                    "prompt_category": global_prompt.prompt_category.name,
                    "prompt": global_prompt.content,
                    "description": global_prompt.description,
                    "arguments": prompt_args,
                    "id": str(global_prompt.id),
                    "source": "global",
                    "is_override": False,
                    "can_revert": False,
                }
                session.close()
                return result

        session.close()
        return None

    # =========================================================================
    # Server-level prompt management (super admin only)
    # =========================================================================

    def get_server_prompts(self, include_internal: bool = False):
        """Get all server-level prompts. Super admin function."""
        session = get_session()
        query = session.query(ServerPrompt)
        if not include_internal:
            query = query.filter(ServerPrompt.is_internal == False)
        prompts = query.all()

        result = []
        for prompt in prompts:
            try:
                prompt_args = [arg.name for arg in prompt.arguments]
            except:
                prompt_args = []
            result.append(
                {
                    "id": str(prompt.id),
                    "name": prompt.name,
                    "category": prompt.category.name,
                    "content": prompt.content,
                    "description": prompt.description,
                    "arguments": prompt_args,
                    "is_internal": prompt.is_internal,
                }
            )

        session.close()
        return result

    def get_server_prompt_by_id(self, prompt_id: str):
        """Get a specific server-level prompt by ID. Super admin function."""
        session = get_session()
        prompt = (
            session.query(ServerPrompt).filter(ServerPrompt.id == prompt_id).first()
        )

        if not prompt:
            session.close()
            return None
        try:
            prompt_args = [arg.name for arg in prompt.arguments]
        except:
            prompt_args = []

        result = {
            "id": str(prompt.id),
            "name": prompt.name,
            "category": prompt.category.name if prompt.category else "Default",
            "content": (
                prompt.content if prompt.name != "Think About It" else THINK_ABOUT_IT
            ),
            "description": prompt.description,
            "arguments": prompt_args,
            "is_internal": prompt.is_internal,
            "created_at": str(prompt.created_at) if prompt.created_at else None,
            "updated_at": str(prompt.updated_at) if prompt.updated_at else None,
        }

        session.close()
        return result

    def add_server_prompt(
        self,
        name: str,
        content: str,
        category: str = "Default",
        description: str = "",
        is_internal: bool = False,
    ):
        """Add a server-level prompt. Super admin function."""
        session = get_session()

        # Get or create category
        cat = (
            session.query(ServerPromptCategory)
            .filter(ServerPromptCategory.name == category)
            .first()
        )
        if not cat:
            cat = ServerPromptCategory(
                name=category,
                description=f"{category} category",
                is_internal=is_internal,
            )
            session.add(cat)
            session.commit()

        prompt = ServerPrompt(
            name=name,
            content=content,
            description=description,
            category_id=cat.id,
            is_internal=is_internal,
        )
        session.add(prompt)
        session.commit()

        # Add arguments
        for arg in self._get_prompt_args_from_content(content):
            session.add(ServerPromptArgument(prompt_id=prompt.id, name=arg))
        session.commit()

        prompt_id = str(prompt.id)
        session.close()
        return prompt_id

    def update_server_prompt(
        self,
        prompt_id: str,
        name: str = None,
        content: str = None,
        description: str = None,
        is_internal: bool = None,
    ):
        """Update a server-level prompt. Super admin function."""
        session = get_session()
        prompt = (
            session.query(ServerPrompt).filter(ServerPrompt.id == prompt_id).first()
        )
        if not prompt:
            session.close()
            raise Exception("Server prompt not found")

        if name is not None:
            prompt.name = name
        if content is not None:
            prompt.content = content
            # Update arguments
            session.query(ServerPromptArgument).filter(
                ServerPromptArgument.prompt_id == prompt_id
            ).delete()
            for arg in self._get_prompt_args_from_content(content):
                session.add(ServerPromptArgument(prompt_id=prompt_id, name=arg))
        if description is not None:
            prompt.description = description
        if is_internal is not None:
            prompt.is_internal = is_internal

        session.commit()
        session.close()

    def delete_server_prompt(self, prompt_id: str):
        """Delete a server-level prompt. Super admin function."""
        session = get_session()
        prompt = (
            session.query(ServerPrompt).filter(ServerPrompt.id == prompt_id).first()
        )
        if not prompt:
            session.close()
            raise Exception("Server prompt not found")
        session.delete(prompt)
        session.commit()
        session.close()

    def get_server_prompt_categories(self, include_internal: bool = False):
        """Get all server-level prompt categories. Super admin function."""
        session = get_session()
        query = session.query(ServerPromptCategory)
        if not include_internal:
            query = query.filter(ServerPromptCategory.is_internal == False)
        categories = query.all()

        result = [
            {
                "id": str(cat.id),
                "name": cat.name,
                "description": cat.description,
                "is_internal": cat.is_internal,
            }
            for cat in categories
        ]
        session.close()
        return result

    def add_server_prompt_category(
        self, name: str, description: str = "", is_internal: bool = False
    ):
        """Add a server-level prompt category. Super admin function."""
        session = get_session()
        cat = ServerPromptCategory(
            name=name, description=description, is_internal=is_internal
        )
        session.add(cat)
        session.commit()
        cat_id = str(cat.id)
        session.close()
        return cat_id

    def delete_server_prompt_category(self, category_id: str):
        """Delete a server-level prompt category. Super admin function."""
        session = get_session()
        cat = (
            session.query(ServerPromptCategory)
            .filter(ServerPromptCategory.id == category_id)
            .first()
        )
        if not cat:
            session.close()
            raise Exception("Category not found")
        session.delete(cat)
        session.commit()
        session.close()

    # =========================================================================
    # Company-level prompt management (company admin only)
    # =========================================================================

    def get_company_prompts(self, company_id: str = None):
        """Get all company-level prompts. Company admin function."""
        target_company = company_id or self.company_id
        if not target_company:
            return []

        session = get_session()
        prompts = (
            session.query(CompanyPrompt)
            .filter(CompanyPrompt.company_id == target_company)
            .all()
        )

        result = []
        for prompt in prompts:
            try:
                prompt_args = [arg.name for arg in prompt.arguments]
            except:
                prompt_args = []
            result.append(
                {
                    "id": str(prompt.id),
                    "name": prompt.name,
                    "category": prompt.category.name,
                    "content": prompt.content,
                    "description": prompt.description,
                    "arguments": prompt_args,
                    "server_prompt_id": (
                        str(prompt.server_prompt_id)
                        if prompt.server_prompt_id
                        else None
                    ),
                }
            )

        session.close()
        return result

    def get_company_prompt_by_id(self, prompt_id: str, company_id: str = None):
        """Get a specific company-level prompt by ID. Company admin function."""
        target_company = company_id or self.company_id
        if not target_company:
            return None

        session = get_session()
        prompt = (
            session.query(CompanyPrompt)
            .filter(
                CompanyPrompt.id == prompt_id,
                CompanyPrompt.company_id == target_company,
            )
            .first()
        )

        if not prompt:
            session.close()
            return None

        try:
            prompt_args = [arg.name for arg in prompt.arguments]
        except:
            prompt_args = []

        result = {
            "id": str(prompt.id),
            "name": prompt.name,
            "category": prompt.category.name if prompt.category else "Default",
            "content": prompt.content,
            "description": prompt.description,
            "arguments": prompt_args,
            "server_prompt_id": (
                str(prompt.server_prompt_id) if prompt.server_prompt_id else None
            ),
            "created_at": str(prompt.created_at) if prompt.created_at else None,
            "updated_at": str(prompt.updated_at) if prompt.updated_at else None,
        }

        session.close()
        return result

    def add_company_prompt(
        self,
        name: str,
        content: str,
        category: str = "Default",
        description: str = "",
        company_id: str = None,
    ):
        """Add a company-level prompt. Company admin function."""
        target_company = company_id or self.company_id
        if not target_company:
            raise Exception("No company context")

        session = get_session()

        # Get or create category
        cat = (
            session.query(CompanyPromptCategory)
            .filter(
                CompanyPromptCategory.company_id == target_company,
                CompanyPromptCategory.name == category,
            )
            .first()
        )
        if not cat:
            cat = CompanyPromptCategory(
                company_id=target_company,
                name=category,
                description=f"{category} category",
            )
            session.add(cat)
            session.commit()

        prompt = CompanyPrompt(
            company_id=target_company,
            name=name,
            content=content,
            description=description,
            category_id=cat.id,
        )
        session.add(prompt)
        session.commit()

        # Add arguments
        for arg in self._get_prompt_args_from_content(content):
            session.add(CompanyPromptArgument(prompt_id=prompt.id, name=arg))
        session.commit()

        prompt_id = str(prompt.id)
        session.close()
        return prompt_id

    def update_company_prompt(
        self,
        prompt_id: str,
        name: str = None,
        content: str = None,
        description: str = None,
    ):
        """Update a company-level prompt. Company admin function."""
        session = get_session()
        prompt = (
            session.query(CompanyPrompt)
            .filter(
                CompanyPrompt.id == prompt_id,
                CompanyPrompt.company_id == self.company_id,
            )
            .first()
        )
        if not prompt:
            session.close()
            raise Exception("Company prompt not found")

        if name is not None:
            prompt.name = name
        if content is not None:
            prompt.content = content
            # Update arguments
            session.query(CompanyPromptArgument).filter(
                CompanyPromptArgument.prompt_id == prompt_id
            ).delete()
            for arg in self._get_prompt_args_from_content(content):
                session.add(CompanyPromptArgument(prompt_id=prompt_id, name=arg))
        if description is not None:
            prompt.description = description

        session.commit()
        session.close()

    def delete_company_prompt(self, prompt_id: str):
        """Delete a company-level prompt. Company admin function."""
        session = get_session()
        prompt = (
            session.query(CompanyPrompt)
            .filter(
                CompanyPrompt.id == prompt_id,
                CompanyPrompt.company_id == self.company_id,
            )
            .first()
        )
        if not prompt:
            session.close()
            raise Exception("Company prompt not found")
        session.query(CompanyPromptArgument).filter(
            CompanyPromptArgument.prompt_id == prompt_id
        ).delete()
        session.delete(prompt)
        session.commit()
        session.close()

    def share_prompt_to_company(self, prompt_id: str):
        """
        Share a user's prompt to the company level.
        Company admin function.
        """
        session = get_session()

        # Get user's prompt
        user_prompt = (
            session.query(Prompt)
            .filter(
                Prompt.id == prompt_id,
                Prompt.user_id == self.user_id,
            )
            .first()
        )
        if not user_prompt:
            session.close()
            raise Exception("Prompt not found")

        if not self.company_id:
            session.close()
            raise Exception("No company context")

        # Create company version
        category = user_prompt.prompt_category.name
        session.close()

        return self.add_company_prompt(
            name=user_prompt.name,
            content=user_prompt.content,
            category=category,
            description=user_prompt.description,
        )

    def get_company_prompt_categories(self, company_id: str = None):
        """Get all company-level prompt categories. Company admin function."""
        target_company = company_id or self.company_id
        if not target_company:
            return []

        session = get_session()
        categories = (
            session.query(CompanyPromptCategory)
            .filter(CompanyPromptCategory.company_id == target_company)
            .all()
        )

        result = [
            {
                "id": str(cat.id),
                "name": cat.name,
                "description": cat.description,
                "server_category_id": (
                    str(cat.server_category_id) if cat.server_category_id else None
                ),
            }
            for cat in categories
        ]
        session.close()
        return result

    def add_company_prompt_category(
        self, name: str, description: str = "", company_id: str = None
    ):
        """Add a company-level prompt category. Company admin function."""
        target_company = company_id or self.company_id
        if not target_company:
            raise Exception("No company context")

        session = get_session()
        cat = CompanyPromptCategory(
            company_id=target_company,
            name=name,
            description=description,
        )
        session.add(cat)
        session.commit()
        cat_id = str(cat.id)
        session.close()
        return cat_id

    def delete_company_prompt_category(self, category_id: str):
        """Delete a company-level prompt category. Company admin function."""
        session = get_session()
        cat = (
            session.query(CompanyPromptCategory)
            .filter(
                CompanyPromptCategory.id == category_id,
                CompanyPromptCategory.company_id == self.company_id,
            )
            .first()
        )
        if not cat:
            session.close()
            raise Exception("Category not found")
        session.delete(cat)
        session.commit()
        session.close()
