#!/usr/bin/env python3
"""
Agent Loop - Tool-using LLM controller for Echo.

WO-ECHO-AGENT-LOOP-01: Iterative loop where LLM calls tools,
observes results, and continues until task complete.

This makes Echo a local coding agent that can iterate on tasks
like Claude Code, but using local models via Ollama.

Architecture:
    User Request → AgentLoop → LLM → Tool Call → Agent Execute → Observe → Repeat
                                  ↓
                              Done? → Final Response

Usage:
    from helix_cdc.echo.agent_loop import AgentLoop, AgentLoopConfig

    loop = AgentLoop(AgentLoopConfig())
    result = await loop.run("find where Se formula is computed and add a docstring")

CLI:
    helix-echo --agent "run tests"
"""

from __future__ import annotations

import asyncio
import hashlib
import json
import os
import time
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional

# Project root
PROJECT_ROOT = Path(__file__).resolve().parents[2]

# Environment configuration
DEFAULT_MODEL = os.environ.get("ECHO_AGENT_MODEL", "qwen2.5-coder:7b")
DEFAULT_MAX_TURNS = int(os.environ.get("ECHO_AGENT_MAX_TURNS", "12"))
DEFAULT_TEMPERATURE = float(os.environ.get("ECHO_AGENT_TEMPERATURE", "0.1"))
DEFAULT_BACKEND = os.environ.get("ECHO_AGENT_BACKEND", "ollama")


# System prompt for agent mode
AGENT_SYSTEM_PROMPT = """You are Echo, a local coding agent for the helix-cdc repository.

You have access to these tools:
- codebase: Search code with ripgrep. Arguments: {"query": "search term"}
- test: Run allowlisted test commands. Arguments: {"query": "pytest tests/ -v"}
- patch: Apply code edits with test gates. Arguments: {"query": "edit description", "patch_spec": {...}}
- plan: Generate work order plans. Arguments: {"query": "what to plan"}
- search: Search local documentation. Arguments: {"query": "search term"}
- chat: Get LLM responses for general questions. Arguments: {"query": "question"}

When given a task:
1. Search the codebase to understand context
2. Make a plan if the task is complex
3. Apply patches one at a time
4. Run tests after each change
5. Iterate until tests pass or task is complete

To use a tool, output a JSON block like this:
```tool_call
{"name": "codebase", "arguments": {"query": "Se formula"}}
```

After receiving tool results, analyze them and either:
- Call another tool if you need more information
- Apply a patch if you're ready to make changes
- Provide a final summary if the task is complete

When done, output:
```tool_call
{"name": "done", "arguments": {"summary": "What was accomplished"}}
```

Important:
- Always search the codebase before making changes
- Run tests after making patches
- Be concise in your responses
- If you encounter errors, try a different approach
"""


@dataclass
class AgentLoopConfig:
    """Configuration for the agent loop."""
    max_turns: int = DEFAULT_MAX_TURNS
    model: str = DEFAULT_MODEL
    temperature: float = DEFAULT_TEMPERATURE
    backend: str = DEFAULT_BACKEND  # "ollama" or "llama_cpp"
    receipt_dir: str = "receipts/agent_loop"
    verbose: bool = False


@dataclass
class ToolExecution:
    """Record of a single tool execution."""
    name: str
    arguments: Dict[str, Any]
    result: Any
    elapsed_ms: float
    success: bool
    error: Optional[str] = None

    def to_dict(self) -> Dict[str, Any]:
        return {
            "name": self.name,
            "arguments": self.arguments,
            "result": self.result if not isinstance(self.result, dict) else self.result,
            "elapsed_ms": self.elapsed_ms,
            "success": self.success,
            "error": self.error,
        }


@dataclass
class AgentLoopResult:
    """Result from running the agent loop."""
    request: str
    response: str
    tool_calls: List[ToolExecution]
    turns: int
    total_elapsed_ms: float
    success: bool
    error: Optional[str] = None
    receipt_path: Optional[str] = None

    def to_dict(self) -> Dict[str, Any]:
        return {
            "request": self.request,
            "response": self.response,
            "tool_calls": [tc.to_dict() for tc in self.tool_calls],
            "turns": self.turns,
            "total_elapsed_ms": self.total_elapsed_ms,
            "success": self.success,
            "error": self.error,
            "receipt_path": self.receipt_path,
        }


class AgentLoop:
    """
    Tool-using LLM controller.

    Implements the iterative loop:
    1. Send prompt to LLM
    2. Parse tool calls from response
    3. Execute tools via dispatchable agents
    4. Append observations to context
    5. Repeat until done or max turns
    """

    def __init__(self, config: Optional[AgentLoopConfig] = None):
        """
        Initialize the agent loop.

        Args:
            config: AgentLoopConfig for behavior customization
        """
        self.config = config or AgentLoopConfig()
        self.receipt_dir = PROJECT_ROOT / self.config.receipt_dir
        self.history: List[Dict[str, Any]] = []
        self.tool_executions: List[ToolExecution] = []
        self._backend = None

    def _get_backend(self):
        """Get or initialize the LLM backend."""
        if self._backend is not None:
            return self._backend

        if self.config.backend == "ollama":
            from echo_cog.backend_ollama import (
                chat_with_tools,
                OllamaConfig,
                check_ollama_health,
            )

            # Check health
            is_healthy, message = check_ollama_health()
            if not is_healthy:
                raise RuntimeError(f"Ollama not available: {message}")

            self._backend = {
                "type": "ollama",
                "chat_with_tools": chat_with_tools,
                "config": OllamaConfig(
                    model=self.config.model,
                    temperature=self.config.temperature,
                ),
            }
        else:
            # Fall back to llama_cpp
            from echo_cog.backend_llama_cpp import (
                run_llama_cpp,
                LlamaCppBackendConfig,
            )

            model_path = os.environ.get("ECHO_MODEL_PATH")
            if not model_path:
                raise RuntimeError("ECHO_MODEL_PATH not set for llama_cpp backend")

            self._backend = {
                "type": "llama_cpp",
                "run": run_llama_cpp,
                "config": LlamaCppBackendConfig(model_path=model_path),
            }

        return self._backend

    def _build_tool_list(self) -> List[Dict[str, Any]]:
        """
        Build OpenAI-compatible tool schemas from registered agents.

        Returns list of tool definitions for the LLM.
        """
        tools = []

        # Import agent registry
        try:
            from helix_cdc.echo.agents import AGENT_REGISTRY
            # Trigger agent registration
            from helix_cdc.echo.agents import (
                codebase_agent,
                test_agent,
                patch_agent,
                plan_agent,
                search_agent,
                chat_agent,
            )
        except ImportError:
            # Return minimal tool set
            pass

        # Define tool schemas for each agent type
        tool_schemas = {
            "codebase": {
                "type": "function",
                "function": {
                    "name": "codebase",
                    "description": "Search the helix-cdc codebase using ripgrep. Use this to find files, functions, classes, or any code patterns.",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "query": {
                                "type": "string",
                                "description": "Search query - can be a regex pattern or plain text"
                            }
                        },
                        "required": ["query"]
                    }
                }
            },
            "test": {
                "type": "function",
                "function": {
                    "name": "test",
                    "description": "Run test commands (pytest, make targets). Commands are validated against an allowlist.",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "query": {
                                "type": "string",
                                "description": "Test command to run, e.g., 'pytest tests/ -v' or 'make test'"
                            }
                        },
                        "required": ["query"]
                    }
                }
            },
            "patch": {
                "type": "function",
                "function": {
                    "name": "patch",
                    "description": "Apply code edits to files. Includes pre/post test gates and automatic rollback.",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "query": {
                                "type": "string",
                                "description": "Description of the edit to make"
                            },
                            "file": {
                                "type": "string",
                                "description": "File path to edit (relative to project root)"
                            },
                            "operation": {
                                "type": "string",
                                "enum": ["add_comment", "append", "prepend", "replace"],
                                "description": "Type of edit operation"
                            },
                            "content": {
                                "type": "string",
                                "description": "Content to add or new content for replace"
                            },
                            "old_text": {
                                "type": "string",
                                "description": "Text to replace (for replace operation)"
                            }
                        },
                        "required": ["query"]
                    }
                }
            },
            "plan": {
                "type": "function",
                "function": {
                    "name": "plan",
                    "description": "Generate a work order plan for complex tasks.",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "query": {
                                "type": "string",
                                "description": "Task to plan"
                            }
                        },
                        "required": ["query"]
                    }
                }
            },
            "search": {
                "type": "function",
                "function": {
                    "name": "search",
                    "description": "Search local documentation and notes.",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "query": {
                                "type": "string",
                                "description": "Search query"
                            }
                        },
                        "required": ["query"]
                    }
                }
            },
            "done": {
                "type": "function",
                "function": {
                    "name": "done",
                    "description": "Signal that the task is complete. Call this when finished.",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "summary": {
                                "type": "string",
                                "description": "Summary of what was accomplished"
                            }
                        },
                        "required": ["summary"]
                    }
                }
            },
        }

        # Only include tools for available agents
        try:
            from helix_cdc.echo.agents import AGENT_REGISTRY, get_agent

            for agent_type, schema in tool_schemas.items():
                if agent_type == "done":
                    # Always include done
                    tools.append(schema)
                else:
                    agent = get_agent(agent_type)
                    if agent and agent.available():
                        tools.append(schema)
        except ImportError:
            # Include all tools if we can't check
            tools = list(tool_schemas.values())

        return tools

    async def _execute_tool(self, name: str, arguments: Dict[str, Any]) -> ToolExecution:
        """
        Execute a tool via the appropriate dispatchable agent.

        Args:
            name: Tool/agent name
            arguments: Arguments for the tool

        Returns:
            ToolExecution with results
        """
        start = time.perf_counter()

        # Handle "done" signal
        if name == "done":
            elapsed_ms = (time.perf_counter() - start) * 1000
            return ToolExecution(
                name=name,
                arguments=arguments,
                result={"done": True, "summary": arguments.get("summary", "")},
                elapsed_ms=elapsed_ms,
                success=True,
            )

        # Get agent
        try:
            from helix_cdc.echo.agents import get_agent

            agent = get_agent(name)
            if agent is None:
                elapsed_ms = (time.perf_counter() - start) * 1000
                return ToolExecution(
                    name=name,
                    arguments=arguments,
                    result=None,
                    elapsed_ms=elapsed_ms,
                    success=False,
                    error=f"Agent '{name}' not found",
                )

            if not agent.available():
                elapsed_ms = (time.perf_counter() - start) * 1000
                return ToolExecution(
                    name=name,
                    arguments=arguments,
                    result=None,
                    elapsed_ms=elapsed_ms,
                    success=False,
                    error=f"Agent '{name}' not available",
                )

            # Build context from arguments
            query = arguments.get("query", json.dumps(arguments))
            context = {
                "agent_loop": True,
                **{k: v for k, v in arguments.items() if k != "query"},
            }

            # For patch agent, build patch_spec
            if name == "patch" and "file" in arguments:
                context["patch_spec"] = {
                    "files": [arguments["file"]],
                    "operation": arguments.get("operation", "replace"),
                    "content": arguments.get("content", ""),
                    "old_text": arguments.get("old_text", ""),
                }

            # Execute
            result = await agent.execute(query, context)

            elapsed_ms = (time.perf_counter() - start) * 1000

            # Convert AgentResult to dict
            result_dict = result.to_dict() if hasattr(result, "to_dict") else result

            return ToolExecution(
                name=name,
                arguments=arguments,
                result=result_dict,
                elapsed_ms=elapsed_ms,
                success=result.error is None,
                error=result.error,
            )

        except Exception as e:
            elapsed_ms = (time.perf_counter() - start) * 1000
            return ToolExecution(
                name=name,
                arguments=arguments,
                result=None,
                elapsed_ms=elapsed_ms,
                success=False,
                error=str(e),
            )

    async def _call_llm(self) -> Dict[str, Any]:
        """
        Call the LLM with current history.

        Returns dict with:
            - content: LLM response text
            - tool_calls: List of parsed tool calls
            - done: Whether LLM signaled completion
            - error: Error message if failed
        """
        backend = self._get_backend()

        if backend["type"] == "ollama":
            from echo_cog.backend_ollama import chat_with_tools

            tools = self._build_tool_list()
            response = chat_with_tools(
                messages=self.history,
                tools=tools,
                config=backend["config"],
            )

            if response.error:
                return {
                    "content": "",
                    "tool_calls": [],
                    "done": False,
                    "error": response.error,
                }

            return {
                "content": response.content,
                "tool_calls": response.tool_calls,
                "done": response.done,
                "error": None,
            }

        else:
            # llama_cpp backend - use markdown parsing
            from echo_cog.backend_ollama import parse_tool_calls_from_markdown

            # Build prompt from history
            prompt_parts = []
            for msg in self.history:
                role = msg["role"]
                content = msg["content"]
                if role == "system":
                    prompt_parts.append(f"System: {content}\n")
                elif role == "user":
                    prompt_parts.append(f"User: {content}\n")
                elif role == "assistant":
                    prompt_parts.append(f"Assistant: {content}\n")
                elif role == "tool":
                    prompt_parts.append(f"Tool Result ({msg.get('name', 'unknown')}): {content}\n")

            prompt_parts.append("Assistant: ")
            full_prompt = "\n".join(prompt_parts)

            response = backend["run"](
                backend["config"],
                full_prompt,
                max_tokens=1024,
            )

            tool_calls = parse_tool_calls_from_markdown(response)
            done = '{"done": true' in response.lower() or not tool_calls

            return {
                "content": response,
                "tool_calls": tool_calls,
                "done": done,
                "error": None,
            }

    def _format_tool_result(self, execution: ToolExecution) -> str:
        """Format tool execution result for the LLM context."""
        if execution.error:
            return f"Error: {execution.error}"

        result = execution.result
        if isinstance(result, dict):
            # Summarize large results
            if "result" in result and isinstance(result["result"], dict):
                inner = result["result"]
                if "matches" in inner:
                    # Codebase search results
                    matches = inner["matches"][:5]  # Limit to 5
                    match_lines = []
                    for m in matches:
                        match_lines.append(f"  {m['file']}:{m['line']}: {m['content'][:80]}")
                    return f"Found {len(inner['matches'])} matches:\n" + "\n".join(match_lines)

                if "status" in inner:
                    # Test/patch results
                    status = inner["status"]
                    preview = inner.get("stdout_preview", inner.get("answer", ""))[:500]
                    return f"Status: {status}\n{preview}"

            # Generic dict
            return json.dumps(result, indent=2, default=str)[:1000]

        return str(result)[:1000]

    def _save_receipt(self, result: AgentLoopResult) -> str:
        """Save execution receipt to disk."""
        self.receipt_dir.mkdir(parents=True, exist_ok=True)

        timestamp = datetime.now(timezone.utc)
        receipt_id = hashlib.sha256(
            f"agent_loop:{timestamp.isoformat()}:{time.time_ns()}".encode()
        ).hexdigest()[:16]

        filename = f"loop_{timestamp.strftime('%Y%m%d')}_{receipt_id}.json"
        receipt_path = self.receipt_dir / filename

        receipt = {
            "receipt_id": receipt_id,
            "timestamp": timestamp.isoformat(),
            "request": result.request,
            "response": result.response,
            "tool_calls": [tc.to_dict() for tc in result.tool_calls],
            "turns": result.turns,
            "total_elapsed_ms": result.total_elapsed_ms,
            "success": result.success,
            "error": result.error,
            "config": {
                "model": self.config.model,
                "backend": self.config.backend,
                "max_turns": self.config.max_turns,
            },
        }

        with open(receipt_path, "w") as f:
            json.dump(receipt, f, indent=2, default=str)

        return str(receipt_path.relative_to(PROJECT_ROOT))

    async def run(self, user_request: str) -> AgentLoopResult:
        """
        Execute the agent loop.

        Args:
            user_request: The user's task/request

        Returns:
            AgentLoopResult with final response and execution history
        """
        start_time = time.perf_counter()

        # Initialize history
        self.history = [
            {"role": "system", "content": AGENT_SYSTEM_PROMPT},
            {"role": "user", "content": user_request},
        ]
        self.tool_executions = []

        if self.config.verbose:
            print(f"[AgentLoop] Starting with request: {user_request[:100]}...")

        final_response = ""

        for turn in range(self.config.max_turns):
            if self.config.verbose:
                print(f"[AgentLoop] Turn {turn + 1}/{self.config.max_turns}")

            # 1. Call LLM
            try:
                response = await self._call_llm()
            except Exception as e:
                elapsed_ms = (time.perf_counter() - start_time) * 1000
                result = AgentLoopResult(
                    request=user_request,
                    response="",
                    tool_calls=self.tool_executions,
                    turns=turn + 1,
                    total_elapsed_ms=elapsed_ms,
                    success=False,
                    error=f"LLM call failed: {e}",
                )
                result.receipt_path = self._save_receipt(result)
                return result

            if response.get("error"):
                elapsed_ms = (time.perf_counter() - start_time) * 1000
                result = AgentLoopResult(
                    request=user_request,
                    response="",
                    tool_calls=self.tool_executions,
                    turns=turn + 1,
                    total_elapsed_ms=elapsed_ms,
                    success=False,
                    error=response["error"],
                )
                result.receipt_path = self._save_receipt(result)
                return result

            content = response.get("content", "")
            tool_calls = response.get("tool_calls", [])

            # Add assistant response to history
            self.history.append({"role": "assistant", "content": content})

            if self.config.verbose:
                print(f"[AgentLoop] LLM response: {content[:200]}...")
                print(f"[AgentLoop] Tool calls: {[tc.get('name') for tc in tool_calls]}")

            # 2. Check for completion
            if response.get("done") or not tool_calls:
                final_response = content
                break

            # 3. Execute tool calls
            for tool_call in tool_calls:
                name = tool_call.get("name", "")
                arguments = tool_call.get("arguments", {})

                # Check for "done" signal
                if name == "done":
                    final_response = arguments.get("summary", content)
                    elapsed_ms = (time.perf_counter() - start_time) * 1000
                    result = AgentLoopResult(
                        request=user_request,
                        response=final_response,
                        tool_calls=self.tool_executions,
                        turns=turn + 1,
                        total_elapsed_ms=elapsed_ms,
                        success=True,
                    )
                    result.receipt_path = self._save_receipt(result)
                    return result

                if self.config.verbose:
                    print(f"[AgentLoop] Executing tool: {name}")

                execution = await self._execute_tool(name, arguments)
                self.tool_executions.append(execution)

                # Add tool result to history
                result_str = self._format_tool_result(execution)
                self.history.append({
                    "role": "tool",
                    "name": name,
                    "content": result_str,
                })

                if self.config.verbose:
                    print(f"[AgentLoop] Tool result: {result_str[:200]}...")

        # Max turns reached or natural completion
        elapsed_ms = (time.perf_counter() - start_time) * 1000

        if not final_response:
            final_response = "Max turns reached without completion"

        result = AgentLoopResult(
            request=user_request,
            response=final_response,
            tool_calls=self.tool_executions,
            turns=min(turn + 1, self.config.max_turns),
            total_elapsed_ms=elapsed_ms,
            success=True if final_response else False,
            error="Max turns exceeded" if turn >= self.config.max_turns - 1 else None,
        )
        result.receipt_path = self._save_receipt(result)
        return result

    def reset(self):
        """Reset the agent loop state."""
        self.history = []
        self.tool_executions = []


def run_agent_loop(request: str, config: Optional[AgentLoopConfig] = None) -> AgentLoopResult:
    """
    Synchronous wrapper for running the agent loop.

    Args:
        request: User's task/request
        config: Optional configuration

    Returns:
        AgentLoopResult
    """
    loop = AgentLoop(config)
    return asyncio.run(loop.run(request))


# CLI interface
if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(
        prog="agent_loop",
        description="Run Echo agent loop on a task",
    )
    parser.add_argument(
        "request",
        nargs="+",
        help="Task to perform",
    )
    parser.add_argument(
        "--model",
        default=DEFAULT_MODEL,
        help=f"Model to use (default: {DEFAULT_MODEL})",
    )
    parser.add_argument(
        "--max-turns",
        type=int,
        default=DEFAULT_MAX_TURNS,
        help=f"Maximum turns (default: {DEFAULT_MAX_TURNS})",
    )
    parser.add_argument(
        "--verbose", "-v",
        action="store_true",
        help="Verbose output",
    )
    parser.add_argument(
        "--json",
        action="store_true",
        help="Output as JSON",
    )

    args = parser.parse_args()

    config = AgentLoopConfig(
        model=args.model,
        max_turns=args.max_turns,
        verbose=args.verbose,
    )

    request = " ".join(args.request)
    result = run_agent_loop(request, config)

    if args.json:
        print(json.dumps(result.to_dict(), indent=2))
    else:
        print(f"\n{'='*60}")
        print(f"Request: {result.request}")
        print(f"{'='*60}")
        print(f"Response:\n{result.response}")
        print(f"{'='*60}")
        print(f"Turns: {result.turns}")
        print(f"Tool calls: {len(result.tool_calls)}")
        print(f"Elapsed: {result.total_elapsed_ms:.1f}ms")
        print(f"Success: {result.success}")
        if result.error:
            print(f"Error: {result.error}")
        if result.receipt_path:
            print(f"Receipt: {result.receipt_path}")
