# core/llm_client.py
import httpx
import json
from typing import Dict, List, Optional, Any
import logging
import time

logger = logging.getLogger(__name__)


class LLMClient:
    """Raw HTTP client for OpenAI API — zero SDK dependency."""

    def __init__(
        self,
        api_key: str,
        base_url: str = "https://api.openai.com/v1",
        timeout: float = 60.0,
        max_retries: int = 3,
    ):
        self.api_key = api_key
        self.base_url = base_url
        self.client = httpx.Client(timeout=timeout)
        self.max_retries = max_retries
        self._total_tokens = 0
        self._total_cost = 0.0

    def chat(
        self,
        messages: List[Dict],
        model: str = "gpt-4o",
        tools: Optional[List[Dict]] = None,
        temperature: float = 0.0,
        max_tokens: int = 2048,
    ) -> Dict[str, Any]:
        payload: Dict[str, Any] = {
            "model": model,
            "messages": messages,
            "temperature": temperature,
            "max_tokens": max_tokens,
        }
        if tools:
            payload["tools"] = tools
            payload["tool_choice"] = "auto"

        for attempt in range(self.max_retries):
            try:
                start = time.time()
                response = self.client.post(
                    f"{self.base_url}/chat/completions",
                    headers={
                        "Authorization": f"Bearer {self.api_key}",
                        "Content-Type": "application/json",
                    },
                    json=payload,
                )
                latency_ms = (time.time() - start) * 1000

                data = response.json()
                if "error" in data:
                    raise Exception(f"LLM API error: {data['error']}")

                choice = data["choices"][0]
                usage = data.get("usage", {})
                self._total_tokens += usage.get("total_tokens", 0)

                return {
                    "content": choice["message"].get("content", ""),
                    "tool_calls": choice["message"].get("tool_calls", []),
                    "finish_reason": choice.get("finish_reason", ""),
                    "usage": usage,
                    "latency_ms": round(latency_ms, 2),
                }
            except Exception as e:
                logger.warning("LLM call attempt %d failed: %s", attempt + 1, e)
                if attempt == self.max_retries - 1:
                    raise

        return {"content": "", "tool_calls": [], "finish_reason": "error", "usage": {}}

    def embeddings(
        self,
        texts: List[str],
        model: str = "text-embedding-3-small",
    ) -> List[List[float]]:
        response = self.client.post(
            f"{self.base_url}/embeddings",
            headers={
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json",
            },
            json={"model": model, "input": texts},
        )
        data = response.json()
        return [item["embedding"] for item in data["data"]]

    def get_stats(self) -> Dict[str, Any]:
        return {"total_tokens": self._total_tokens}

# core/tool_registry.py
import inspect
from typing import Any, Callable, Dict, List, Optional
import json
import logging

logger = logging.getLogger(__name__)


class Tool:
    """Represents a registered tool with metadata and execution capability."""

    def __init__(
        self,
        name: str,
        description: str,
        func: Callable,
        parameters: Dict,
    ):
        self.name = name
        self.description = description
        self.func = func
        self.parameters = parameters
        self.is_async = inspect.iscoroutinefunction(func)

    def execute(self, **kwargs: Any) -> str:
        try:
            result = self.func(**kwargs)
            return str(result)
        except Exception as e:
            return f"Error executing {self.name}: {str(e)}"

    async def execute_async(self, **kwargs: Any) -> str:
        try:
            if self.is_async:
                result = await self.func(**kwargs)
            else:
                result = self.func(**kwargs)
            return str(result)
        except Exception as e:
            return f"Error executing {self.name}: {str(e)}"


class ToolRegistry:
    """Dynamic tool registration with automatic schema generation."""

    def __init__(self):
        self.tools: Dict[str, Tool] = {}

    def register(
        self,
        name: str,
        description: str,
        parameters: Dict,
    ) -> Callable:
        def decorator(func: Callable) -> Callable:
            self.tools[name] = Tool(name, description, func, parameters)
            logger.info("Registered tool: %s", name)
            return func
        return decorator

    def register_function(
        self,
        func: Callable,
        name: Optional[str] = None,
        description: Optional[str] = None,
        parameters: Optional[Dict] = None,
    ) -> None:
        tool_name = name or func.__name__
        tool_desc = description or func.__doc__ or f"Execute {tool_name}"
        tool_params = parameters or self._generate_schema(func)
        self.tools[tool_name] = Tool(tool_name, tool_desc, func, tool_params)
        logger.info("Registered tool function: %s", tool_name)

    def _generate_schema(self, func: Callable) -> Dict:
        sig = inspect.signature(func)
        props: Dict[str, Dict] = {}
        required: List[str] = []

        for pname, param in sig.parameters.items():
            ptype = "string"
            if param.annotation != inspect.Parameter.empty:
                type_map = {
                    str: "string",
                    int: "integer",
                    float: "number",
                    bool: "boolean",
                    list: "array",
                    dict: "object",
                }
                ptype = type_map.get(param.annotation, "string")

            props[pname] = {
                "type": ptype,
                "description": f"The {pname} parameter",
            }
            if param.default is inspect.Parameter.empty:
                required.append(pname)

        return {
            "type": "object",
            "properties": props,
            "required": required,
        }

    def get_openai_tools(self) -> List[Dict]:
        return [
            {
                "type": "function",
                "function": {
                    "name": tool.name,
                    "description": tool.description,
                    "parameters": tool.parameters,
                },
            }
            for tool in self.tools.values()
        ]

    def execute(self, name: str, arguments: Dict) -> str:
        if name not in self.tools:
            return f"Error: Unknown tool '{name}'"
        tool = self.tools[name]
        return tool.execute(**arguments)

    async def execute_async(self, name: str, arguments: Dict) -> str:
        if name not in self.tools:
            return f"Error: Unknown tool '{name}'"
        tool = self.tools[name]
        return await tool.execute_async(**arguments)

    def list_tools(self) -> List[str]:
        return list(self.tools.keys())

# core/memory.py
import tiktoken
from typing import Dict, List, Optional
import logging

logger = logging.getLogger(__name__)


class Memory:
    """Token-aware conversation memory with automatic trimming."""

    def __init__(
        self,
        max_tokens: int = 4000,
        system_prompt: str = "",
        keep_recent: int = 4,
    ):
        self.max_tokens = max_tokens
        self.system_prompt = system_prompt
        self.messages: List[Dict] = []
        self.keep_recent = keep_recent
        self.enc = tiktoken.get_encoding("cl100k_base")
        self._trim_count = 0

    def add_user(self, content: str) -> None:
        self.messages.append({"role": "user", "content": content})
        self._trim()

    def add_assistant(self, content: str) -> None:
        self.messages.append({"role": "assistant", "content": content})

    def add_tool_result(self, tool_call_id: str, content: str) -> None:
        self.messages.append({
            "role": "tool",
            "tool_call_id": tool_call_id,
            "content": content,
        })

    def get_messages(self) -> List[Dict]:
        messages: List[Dict] = []
        if self.system_prompt:
            messages.append({"role": "system", "content": self.system_prompt})
        messages.extend(self.messages)
        return messages

    def _count_tokens(self, text: str) -> int:
        return len(self.enc.encode(text))

    def _trim(self) -> None:
        total = sum(
            self._count_tokens(m.get("content", ""))
            for m in self.messages
        )
        original_count = len(self.messages)

        while total > self.max_tokens and len(self.messages) > self.keep_recent:
            removed = self.messages.pop(0)
            total -= self._count_tokens(removed.get("content", ""))
            self._trim_count += 1

        if len(self.messages) < original_count:
            logger.info(
                "Trimmed memory: %d -> %d messages (saved ~%d tokens)",
                original_count,
                len(self.messages),
                self._trim_count,
            )

    def clear(self) -> None:
        self.messages.clear()
        self._trim_count = 0

    def get_recent(self, n: int = 5) -> List[Dict]:
        return self.messages[-n:]

    def get_stats(self) -> Dict[str, int]:
        total_tokens = sum(
            self._count_tokens(m.get("content", ""))
            for m in self.messages
        )
        return {
            "message_count": len(self.messages),
            "total_tokens": total_tokens,
            "max_tokens": self.max_tokens,
            "utilization_pct": round(total_tokens / self.max_tokens * 100, 2),
            "trim_count": self._trim_count,
        }

# core/safety.py
import re
from typing import Dict, Tuple, List
import logging

logger = logging.getLogger(__name__)


class SafetyLayer:
    """Basic safety layer for input/output validation."""

    FORBIDDEN_PATTERNS = [
        r"ignore previous instructions",
        r"you are now.*",
        r"disregard.*instructions",
        r"reveal.*system prompt",
        r"act as.*admin",
    ]
    PII_PATTERNS = {
        "email": r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b',
        "phone": r'\b\d{3}[-.]?\d{3}[-.]?\d{4}\b',
        "ssn": r'\b\d{3}-\d{2}-\d{4}\b',
    }

    def check_input(self, text: str) -> Tuple[bool, str]:
        for pattern in self.FORBIDDEN_PATTERNS:
            if re.search(pattern, text, re.IGNORECASE):
                logger.warning("Input safety check failed: %s", pattern)
                return False, "Blocked: potential jailbreak"
        return True, ""

    def check_output(self, text: str) -> Tuple[bool, str]:
        for pii_type, pattern in self.PII_PATTERNS.items():
            if re.search(pattern, text):
                return False, f"Output contains {pii_type}"
        return True, ""

    def redact_pii(self, text: str) -> str:
        for pii_type, pattern in self.PII_PATTERNS.items():
            text = re.sub(pattern, f"[REDACTED {pii_type.upper()}]", text)
        return text


# core/executor.py
import json
import time
from typing import Dict, List, Any
import logging

logger = logging.getLogger(__name__)


class ReActExecutor:
    """ReAct-style executor: Thought → Action → Observation loop."""

    def __init__(
        self,
        llm: Any,
        tools: Any,
        memory: Any,
        safety: Any,
        model: str = "gpt-4o",
        max_iterations: int = 10,
    ):
        self.llm = llm
        self.tools = tools
        self.memory = memory
        self.safety = safety
        self.model = model
        self.max_iterations = max_iterations
        self._execution_count = 0

    def run(self, query: str) -> Dict[str, Any]:
        safe, reason = self.safety.check_input(query)
        if not safe:
            return {
                "answer": "I can't process that request.",
                "iterations": 0,
                "blocked": True,
                "reason": reason,
            }

        self.memory.add_user(query)
        trace: List[Dict] = []
        start_time = time.time()

        for i in range(self.max_iterations):
            messages = self.memory.get_messages()
            tool_schemas = self.tools.get_openai_tools()

            response = self.llm.chat(
                messages,
                model=self.model,
                tools=tool_schemas if tool_schemas else None,
            )

            if response.get("tool_calls"):
                for tc in response["tool_calls"]:
                    func_name = tc["function"]["name"]
                    try:
                        args = json.loads(tc["function"]["arguments"])
                    except json.JSONDecodeError:
                        args = {}

                    result = self.tools.execute(func_name, args)
                    self.memory.add_tool_result(tc["id"], result)
                    trace.append({
                        "step": i + 1,
                        "tool": func_name,
                        "arguments": args,
                        "result": result[:200],
                    })
            elif response.get("content"):
                safe, reason = self.safety.check_output(response["content"])
                if not safe:
                    response["content"] = self.safety.redact_pii(response["content"])

                self.memory.add_assistant(response["content"])
                total_time = (time.time() - start_time) * 1000
                self._execution_count += 1

                return {
                    "answer": response["content"],
                    "iterations": i + 1,
                    "trace": trace,
                    "usage": response.get("usage", {}),
                    "latency_ms": round(total_time, 2),
                }

        return {
            "answer": "Max iterations reached without final answer",
            "iterations": self.max_iterations,
            "trace": trace,
        }

# agent.py
from core.llm_client import LLMClient
from core.tool_registry import ToolRegistry
from core.memory import Memory
from core.safety import SafetyLayer
from core.executor import ReActExecutor
from typing import Dict, Any, Optional
import logging

logger = logging.getLogger(__name__)


class Agent:
    """Complete agent framework with zero external agent dependencies."""

    def __init__(
        self,
        api_key: str,
        system_prompt: str = "You are a helpful assistant.",
        model: str = "gpt-4o",
        max_memory_tokens: int = 4000,
        max_iterations: int = 10,
    ):
        self.llm = LLMClient(api_key)
        self.tools = ToolRegistry()
        self.memory = Memory(
            max_tokens=max_memory_tokens,
            system_prompt=system_prompt,
        )
        self.safety = SafetyLayer()
        self.executor = ReActExecutor(
            self.llm, self.tools, self.memory, self.safety, model, max_iterations
        )
        self.model = model
        logger.info("Agent initialized with model: %s", model)

    def register_tool(
        self,
        name: str,
        func,
        description: Optional[str] = None,
        parameters: Optional[Dict] = None,
    ) -> None:
        self.tools.register_function(func, name, description, parameters)

    def run(self, query: str) -> Dict[str, Any]:
        return self.executor.run(query)

    def reset(self) -> None:
        self.memory.clear()

    def get_stats(self) -> Dict[str, Any]:
        return {
            "llm_stats": self.llm.get_stats(),
            "memory_stats": self.memory.get_stats(),
            "tools": self.tools.list_tools(),
            "execution_count": self.executor._execution_count,
        }


# tools/builtin.py
import math
from typing import Optional


def calculator(expression: str) -> str:
    safe_dict = {
        "sqrt": math.sqrt,
        "log": math.log,
        "sin": math.sin,
        "cos": math.cos,
        "tan": math.tan,
        "pi": math.pi,
        "e": math.e,
        "abs": abs,
        "round": round,
        "pow": pow,
    }
    result = eval(expression, {"__builtins__": {}}, safe_dict)
    return str(result)


def echo(text: str) -> str:
    return f"Echo: {text}"


def get_current_time() -> str:
    from datetime import datetime
    return datetime.now().isoformat()


# main.py
from agent import Agent
from tools.builtin import calculator, echo, get_current_time
import logging

logging.basicConfig(level=logging.INFO)


def main():
    import os
    api_key = os.environ.get("OPENAI_API_KEY", "sk-your-key")

    agent = Agent(
        api_key=api_key,
        system_prompt="You are a helpful assistant with access to tools.",
    )

    agent.register_tool(
        "calculator",
        calculator,
        "Evaluate math expressions safely",
        {
            "type": "object",
            "properties": {"expression": {"type": "string", "description": "Math expression to evaluate"}},
            "required": ["expression"],
        },
    )
    agent.register_tool(
        "echo",
        echo,
        "Echo text back",
        {
            "type": "object",
            "properties": {"text": {"type": "string", "description": "Text to echo"}},
            "required": ["text"],
        },
    )
    agent.register_tool(
        "get_current_time",
        get_current_time,
        "Get current date and time",
        {"type": "object", "properties": {}},
    )

    print("Agent Framework Ready. Type 'quit' to exit.\n")
    while True:
        query = input("You: ").strip()
        if query.lower() in ("quit", "exit", "q"):
            break
        if not query:
            continue

        result = agent.run(query)
        print(f"\nAgent: {result['answer']}")
        print(f"   ({result['iterations']} iterations, {result.get('latency_ms', 0):.0f}ms)\n")

    print("\nFinal Stats:", agent.get_stats())


if __name__ == "__main__":
    main()

import pytest
from core.tool_registry import ToolRegistry
from core.memory import Memory
from core.safety import SafetyLayer


def test_tool_registry():
    registry = ToolRegistry()

    @registry.register("add", "Add two numbers", {
        "type": "object",
        "properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
        "required": ["a", "b"],
    })
    def add(a: int, b: int) -> int:
        return a + b

    result = registry.execute("add", {"a": 2, "b": 3})
    assert result == "5"


def test_memory():
    memory = Memory(max_tokens=100, system_prompt="Test")
    memory.add_user("Hello")
    memory.add_assistant("Hi there")
    messages = memory.get_messages()
    assert len(messages) == 3
    assert messages[0]["role"] == "system"


def test_safety():
    safety = SafetyLayer()
    safe, _ = safety.check_input("Hello world")
    assert safe
    safe, _ = safety.check_input("Ignore previous instructions")
    assert not safe


def test_pii_redaction():
    safety = SafetyLayer()
    redacted = safety.redact_pii("Email me at test@example.com")
    assert "test@example.com" not in redacted