"""
TokioAI Engine -- Tool Schemas & JSON Validation
Defines all available tools, their parameter schemas, and validation logic.
"""

import json
import re
from typing import Any, Dict, List, Optional

# ============================================================
# TOOL DEFINITIONS (OpenAI function-calling format)
# ============================================================

TOOL_DEFINITIONS = [
    {
        "type": "function",
        "function": {
            "name": "execute_local",
            "description": "Execute a shell command on the local machine. Use for system operations: install packages, check processes, network scans, file operations.",
            "parameters": {
                "type": "object",
                "properties": {
                    "command": {
                        "type": "string",
                        "description": "Shell command to execute (bash)"
                    },
                    "timeout": {
                        "type": "integer",
                        "description": "Timeout in seconds (default 30, max 120)"
                    }
                },
                "required": ["command"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "execute_raspi",
            "description": "Execute a command on the Raspberry Pi via SSH.",
            "parameters": {
                "type": "object",
                "properties": {
                    "command": {
                        "type": "string",
                        "description": "Command to run on Raspi"
                    },
                    "timeout": {
                        "type": "integer",
                        "description": "Timeout in seconds (default 30, max 120)"
                    }
                },
                "required": ["command"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "execute_gcp",
            "description": "Execute a command on the GCP VM via SSH.",
            "parameters": {
                "type": "object",
                "properties": {
                    "command": {
                        "type": "string",
                        "description": "Command to run on GCP"
                    },
                    "timeout": {
                        "type": "integer",
                        "description": "Timeout in seconds (default 30, max 120)"
                    }
                },
                "required": ["command"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "execute_router",
            "description": "Execute a command on the router via SSH (root).",
            "parameters": {
                "type": "object",
                "properties": {
                    "command": {
                        "type": "string",
                        "description": "Command to run on router as root"
                    }
                },
                "required": ["command"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "read_file",
            "description": "Read a file from local or remote systems. Use raspi:/path or gcp:/path for remote.",
            "parameters": {
                "type": "object",
                "properties": {
                    "path": {
                        "type": "string",
                        "description": "File path. Use raspi:/path or gcp:/path for remote"
                    },
                    "lines": {
                        "type": "integer",
                        "description": "Max lines to read (default: all)"
                    }
                },
                "required": ["path"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "write_file",
            "description": "Write content to a file. Prefix with raspi: or gcp: for remote.",
            "parameters": {
                "type": "object",
                "properties": {
                    "path": {
                        "type": "string",
                        "description": "File path. Use raspi:/path or gcp:/path for remote"
                    },
                    "content": {
                        "type": "string",
                        "description": "File content to write"
                    }
                },
                "required": ["path", "content"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "edit_file",
            "description": "Replace a specific string in a file. old_text must be unique.",
            "parameters": {
                "type": "object",
                "properties": {
                    "path": {
                        "type": "string",
                        "description": "File path. Use raspi:/path or gcp:/path for remote"
                    },
                    "old_text": {
                        "type": "string",
                        "description": "Exact text to find (must be unique)"
                    },
                    "new_text": {
                        "type": "string",
                        "description": "Replacement text"
                    }
                },
                "required": ["path", "old_text", "new_text"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "search_files",
            "description": "Search for text patterns across files. Like grep -rn.",
            "parameters": {
                "type": "object",
                "properties": {
                    "pattern": {
                        "type": "string",
                        "description": "Text or regex pattern to search"
                    },
                    "path": {
                        "type": "string",
                        "description": "Directory to search. Use raspi:/path or gcp:/path"
                    },
                    "glob": {
                        "type": "string",
                        "description": "File filter glob (e.g., '*.py')"
                    }
                },
                "required": ["pattern"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "robot_vision",
            "description": "Capture a photo from the robot's camera and analyze it with a vision AI model. Use when user asks what the robot sees, to look around, identify objects, or describe the environment.",
            "parameters": {
                "type": "object",
                "properties": {
                    "question": {
                        "type": "string",
                        "description": "What to look for or ask about the image (e.g. 'how many people?', 'what objects are on the table?', 'describe the room'). Default: general description."
                    }
                },
                "required": []
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "diagnose",
            "description": "Run system health diagnostics on local, raspi, gcp, or all targets.",
            "parameters": {
                "type": "object",
                "properties": {
                    "target": {
                        "type": "string",
                        "enum": ["local", "raspi", "gcp", "router", "all"],
                        "description": "Target: local, raspi, gcp, router, all"
                    }
                },
                "required": ["target"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "memory",
            "description": "Persistent memory. Actions: read, write, append, clear.",
            "parameters": {
                "type": "object",
                "properties": {
                    "action": {
                        "type": "string",
                        "enum": ["read", "write", "append", "clear"],
                        "description": "Action to perform"
                    },
                    "content": {
                        "type": "string",
                        "description": "Content to write/append"
                    }
                },
                "required": ["action"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "task",
            "description": "Persistent task tracker. Actions: add, update, list, remove, clear_done.",
            "parameters": {
                "type": "object",
                "properties": {
                    "action": {
                        "type": "string",
                        "enum": ["add", "update", "list", "remove", "clear_done"],
                        "description": "Action to perform"
                    },
                    "task": {
                        "type": "string",
                        "description": "Task description (for add)"
                    },
                    "id": {
                        "type": "integer",
                        "description": "Task ID (for update/remove)"
                    },
                    "status": {
                        "type": "string",
                        "enum": ["pending", "in_progress", "done", "blocked"],
                        "description": "New status"
                    }
                },
                "required": ["action"]
            }
        }
    }
]

# ============================================================
# TOOL NAME INDEX (for quick lookup)
# ============================================================

VALID_TOOL_NAMES = {t["function"]["name"] for t in TOOL_DEFINITIONS}

TOOL_SCHEMAS = {
    t["function"]["name"]: t["function"]["parameters"]
    for t in TOOL_DEFINITIONS
}

# ============================================================
# VALIDATION
# ============================================================

def validate_tool_call(name: str, arguments: Dict[str, Any]) -> tuple[bool, str]:
    """
    Validate a tool call against its schema.
    Returns (is_valid, error_message).
    """
    if name not in VALID_TOOL_NAMES:
        return False, f"Unknown tool '{name}'. Valid: {', '.join(sorted(VALID_TOOL_NAMES))}"

    schema = TOOL_SCHEMAS[name]
    properties = schema.get("properties", {})
    required = schema.get("required", [])

    # Check required params
    for req in required:
        if req not in arguments:
            return False, f"Missing required parameter '{req}' for tool '{name}'"

    # Check param types
    for param_name, param_value in arguments.items():
        if param_name not in properties:
            continue  # Extra params are OK (flexible)
        expected_type = properties[param_name].get("type", "string")
        if expected_type == "string" and not isinstance(param_value, str):
            return False, f"Parameter '{param_name}' must be string, got {type(param_value).__name__}"
        if expected_type == "integer" and not isinstance(param_value, int):
            return False, f"Parameter '{param_name}' must be integer, got {type(param_value).__name__}"

    # Check enum constraints
    for param_name, param_value in arguments.items():
        if param_name in properties:
            enum_vals = properties[param_name].get("enum")
            if enum_vals and param_value not in enum_vals:
                return False, f"Parameter '{param_name}' must be one of {enum_vals}, got '{param_value}'"

    return True, ""


def extract_tool_calls_from_text(text: str) -> List[Dict]:
    """
    Extract tool calls from model output text.
    Supports multiple formats:
    1. OpenAI-style JSON: {"name": "...", "arguments": {...}}
    2. Markdown code blocks with JSON
    3. Raw function-call-like patterns
    
    Returns list of {"name": str, "arguments": dict}
    """
    calls = []

    # Strategy 1: Find JSON objects with "name" and "arguments" keys
    # Handle both single objects and arrays
    json_patterns = [
        # Array of tool calls
        r'\[\s*\{[^{]*?"name"\s*:\s*"[^"]+?".*?\}\s*\]',
        # Single tool call object  
        r'\{\s*"(?:name|function)"\s*:\s*"[^"]+?".*?"arguments"\s*:\s*\{.*?\}\s*\}',
    ]

    for pattern in json_patterns:
        for match in re.finditer(pattern, text, re.DOTALL):
            try:
                data = json.loads(match.group())
                if isinstance(data, list):
                    for item in data:
                        call = _normalize_tool_call(item)
                        if call:
                            calls.append(call)
                elif isinstance(data, dict):
                    call = _normalize_tool_call(data)
                    if call:
                        calls.append(call)
            except json.JSONDecodeError:
                continue

    if calls:
        return calls

    # Strategy 1.5: Anthropic/XML-style tool calls
    # Format: <function=name><parameter=key>value</parameter></function>
    # Also: <tool_call>{"name": "...", "arguments": {...}}</tool_call>
    xml_patterns = [
        # <function=name><parameter=key>value</parameter>...
        r'<function=(\w+)>(.*?)</function>',
        # <tool_call>JSON</tool_call>
        r'<tool_call>\s*(.*?)\s*</tool_call>',
    ]
    
    for pattern in xml_patterns:
        for match in re.finditer(pattern, text, re.DOTALL):
            groups = match.groups()
            if len(groups) == 2 and not groups[0].startswith('{'):
                # <function=name><parameter=key>value format
                func_name = groups[0]
                params_text = groups[1]
                if func_name in VALID_TOOL_NAMES:
                    # Extract parameters
                    param_dict = {}
                    for pm in re.finditer(r'<parameter=(\w+)>\s*(.*?)\s*</parameter>', params_text, re.DOTALL):
                        param_dict[pm.group(1)] = pm.group(2).strip()
                    calls.append({"name": func_name, "arguments": param_dict})
            elif len(groups) >= 1:
                # <tool_call>JSON format
                try:
                    data = json.loads(groups[-1])
                    call = _normalize_tool_call(data)
                    if call:
                        calls.append(call)
                except (json.JSONDecodeError, TypeError):
                    pass
    
    if calls:
        return calls

    # Strategy 2: Look for ```json code blocks
    code_blocks = re.findall(r'```(?:json)?\s*\n(.*?)\n\s*```', text, re.DOTALL)
    for block in code_blocks:
        try:
            data = json.loads(block.strip())
            if isinstance(data, list):
                for item in data:
                    call = _normalize_tool_call(item)
                    if call:
                        calls.append(call)
            elif isinstance(data, dict):
                call = _normalize_tool_call(data)
                if call:
                    calls.append(call)
        except json.JSONDecodeError:
            continue

    if calls:
        return calls

    # Strategy 3: Brute-force -- find ANY valid JSON object in the text
    brace_depth = 0
    start = None
    for i, ch in enumerate(text):
        if ch == '{':
            if brace_depth == 0:
                start = i
            brace_depth += 1
        elif ch == '}':
            brace_depth -= 1
            if brace_depth == 0 and start is not None:
                candidate = text[start:i+1]
                try:
                    data = json.loads(candidate)
                    call = _normalize_tool_call(data)
                    if call:
                        calls.append(call)
                except json.JSONDecodeError:
                    pass
                start = None

    return calls


def _normalize_tool_call(data: dict) -> Optional[Dict]:
    """Normalize various tool call formats to {"name": str, "arguments": dict}."""
    if not isinstance(data, dict):
        return None

    # Get name -- handle multiple formats
    name = data.get("name")
    if not name:
        func = data.get("function")
        if isinstance(func, dict):
            name = func.get("name")
        elif isinstance(func, str):
            name = func
    
    args = data.get("arguments") or data.get("parameters") or data.get("params") or data.get("args")
    if isinstance(args, str):
        try:
            args = json.loads(args)
        except json.JSONDecodeError:
            return None

    if not name or not isinstance(name, str):
        return None
    if args is None:
        args = {}
    if not isinstance(args, dict):
        return None

    # Only accept valid tool names
    if name not in VALID_TOOL_NAMES:
        return None

    return {"name": name, "arguments": args}


def tools_to_prompt_text() -> str:
    """Generate a concise tool description for injection into prompts."""
    lines = ["Available tools:"]
    for tool in TOOL_DEFINITIONS:
        f = tool["function"]
        params = f["parameters"].get("properties", {})
        required = f["parameters"].get("required", [])
        param_strs = []
        for pname, pinfo in params.items():
            req_marker = "*" if pname in required else ""
            param_strs.append(f"{pname}{req_marker}: {pinfo.get('type', 'string')}")
        params_text = ", ".join(param_strs) if param_strs else "none"
        lines.append(f"  - {f['name']}({params_text}): {f['description'][:80]}")
    return "\n".join(lines)