# utilities/prompt_compiler.py

import subprocess
import json
import yaml
import re

PROMPT_TEMPLATES = {
    "summarize": "Summarize in one sentence:\n\"\"\"\n{input}\n\"\"\"",
    "mutate_code": "Suggest improvements to this code function, return only modified logic if any:\n\n```python\n{code}\n```",
    "validate_drift": "Compare this summary:\n{summary}\n\nAgainst this input:\n{input}\n\nIs context lost? Reply YES or NO followed by brief justification.",
    "reasoning": "Given this scenario:\n{scenario}\n\nWhat should the agent do next? Reply concisely.",
    "basic_prompt": "You are an agent running on the HiveOS platform. This is the user's input prompt:\n\n{prompt}\n\n",
    "generate_task": (
        "You are a HiveOS automation agent.\n\n"
        "Your job is to translate the following user goal into a valid HiveOS YAML task definition.\n\n"
        "USER GOAL:\n"
        "\"\"\"\n{goal}\n\"\"\"\n\n"
        "Follow the official HiveOS YAML schema strictly. Your output MUST be wrapped in a valid code block: ```yaml ... ```\n\n"
        "Top-level fields:\n"
        "- task_name: A short descriptive name\n"
        "- execution_mode: One of [concurrent, sequential]\n"
        "- capability_required: A list of required agent capabilities\n"
        "- chunks: A flat list of task chunks\n\n"
        "Each chunk must contain:\n"
        "- chunk_id: A unique identifier like 'chunk_1'\n"
        "- required_capability: A single string capability\n"
        "- payload:\n"
        "    - intent: A single-word intent (e.g., render_text, summarize)\n"
        "    - parameters: A dictionary (can be empty)\n\n"
        "Example format:\n"
        "```yaml\n"
        "task_name: Example Task\n"
        "execution_mode: concurrent\n"
        "capability_required:\n"
        "  - display\n"
        "chunks:\n"
        "  - chunk_id: chunk_1\n"
        "    required_capability: display\n"
        "    payload:\n"
        "      intent: render_text\n"
        "      parameters:\n"
        "        text: \"Hello, world!\"\n"
        "```\n\n"
        "Only output valid YAML inside a single code block. Do not include explanations, comments, or markdown."
    )

}

def compile_prompt(intent: str, context: dict) -> str:
    """
    Returns a formatted prompt string given an intent and context dict.
    """
    if intent not in PROMPT_TEMPLATES:
        raise ValueError(f"Unknown intent: {intent}")
    return PROMPT_TEMPLATES[intent].format(**context)

def llm_infer(model_name: str, prompt: str) -> str:
    """
    Runs a one-shot inference using Ollama.
    """
    result = subprocess.run(
        ["ollama", "run", model_name, prompt],
        capture_output=True,
        text=True,
        encoding="utf-8"
    )
    return result.stdout.strip()

def sanitize_output(output: str, expected_type: str = "string"):
    """
    Sanitize LLM output to return clean data.
    """
    if expected_type == "json":
        try:
            match = re.search(r"```json\s*([\s\S]+?)\s*```", output)
            if match:
                return json.loads(match.group(1).strip())
        except Exception:
            return None

    elif expected_type == "yaml":
        try:
            # Try to match a full fenced YAML block
            match = re.search(r"```yaml\s*([\s\S]+?)\s*```", output)
            if match:
                return match.group(1).strip()

            # Fallback if the closing ``` is missing
            if "```yaml" in output:
                return output.split("```yaml", 1)[1].strip()
        except Exception as e:
            print(f"[sanitize_output] YAML parse error: {e}")
            return None

    elif expected_type == "string":
        lines = output.strip().split("\n")
        for line in lines:
            if line.strip() and not line.strip().lower().startswith("sure"):
                return line.strip()

    # Fallback
    return output.strip()[:256]


"""
# Example usage
if __name__ == "__main__":
    context = {"input": "The HiveOS platform coordinates hardware and AI agents using chunked task flows."}
    prompt = compile_prompt("summarize", context)
    result = llm_infer("phi3", prompt)
    clean = sanitize_output(result)
    print(clean)
"""