#!/usr/bin/env python3
"""
Gate 3: LLM Tool-Use Integration Test

WO-ECHO-FIBPI3D-LOOP-01: Tests whether an LLM (Mistral) can:
1. Emit valid tool calls for spatial queries
2. Receive tool results
3. Use results in final answers (not hallucinate)

REQUIRES: GATE3_LLM_MODEL environment variable pointing to a GGUF model.
Without this env var, tests are skipped.

Example:
    GATE3_LLM_MODEL=/path/to/mistral-7b-instruct.gguf pytest tests/test_llm_tool_use.py -v

Uses llama-cpp-python with Mistral 7B Instruct.
"""

import json
import os
import re
import sys
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import pytest
import requests

# Add project root
sys.path.insert(0, str(Path(__file__).parent.parent))

# Configuration
MODEL_PATH = os.environ.get("GATE3_LLM_MODEL")
ENDPOINT = os.environ.get("SPATIAL_ENDPOINT", "http://127.0.0.1:8765")


# Skip all tests if GATE3_LLM_MODEL not set
pytestmark = pytest.mark.skipif(
    not MODEL_PATH,
    reason="LLM tool-use test requires GATE3_LLM_MODEL env var pointing to GGUF model"
)


# Spatial tool schema for Mistral
SPATIAL_TOOL_SCHEMA = {
    "type": "function",
    "function": {
        "name": "spatial_query",
        "description": "Query spatial relationships between objects in a scene. Use this to find distances, nearest neighbors, or check spatial predicates like 'near', 'between', 'left_of', 'above'.",
        "parameters": {
            "type": "object",
            "properties": {
                "query_type": {
                    "type": "string",
                    "enum": ["distance", "k_nearest", "near", "between", "left_of", "above"],
                    "description": "Type of spatial query"
                },
                "a": {
                    "type": "string",
                    "description": "First object name (e.g., 'obj_0007')"
                },
                "b": {
                    "type": "string",
                    "description": "Second object name (for distance, near, left_of, above, between)"
                },
                "c": {
                    "type": "string",
                    "description": "Third object name (only for 'between' query)"
                },
                "k": {
                    "type": "integer",
                    "description": "Number of neighbors (for k_nearest), default 5"
                },
                "threshold": {
                    "type": "number",
                    "description": "Distance threshold (for 'near' query), default 5.0"
                }
            },
            "required": ["query_type", "a"]
        }
    }
}

# System prompt that teaches Mistral to use the tool
SYSTEM_PROMPT = """You are a spatial reasoning assistant. You have access to a spatial_query tool that can answer questions about object positions in a 3D scene.

Available objects are named obj_0000, obj_0001, obj_0002, ... up to obj_0099.

When asked about spatial relationships, you MUST use the spatial_query tool. Do not guess or make up positions.

To use the tool, output a JSON block like this:
```tool_call
{"name": "spatial_query", "arguments": {"query_type": "k_nearest", "a": "obj_0007", "k": 5}}
```

After receiving tool results, provide your final answer based on the actual data.

Query types available:
- distance: Get exact distance between two objects
- k_nearest: Find the k closest objects to a target
- near: Check if two objects are within a threshold distance
- between: Check if object B is geometrically between A and C
- left_of: Check if A.x < B.x
- above: Check if A.z > B.z"""


@dataclass
class ToolCall:
    """Parsed tool call from LLM output."""
    name: str
    arguments: Dict[str, Any]


@dataclass
class Gate3Result:
    """Result of a single Gate 3 test."""
    prompt: str
    llm_output: str
    tool_calls: List[ToolCall]
    tool_results: List[Dict[str, Any]]
    final_answer: str
    passed: bool
    failure_reason: Optional[str] = None


def parse_tool_calls(text: str) -> List[ToolCall]:
    """Extract tool calls from LLM output."""
    calls = []

    # Look for ```tool_call ... ``` blocks
    pattern = r'```tool_call\s*\n?(.*?)\n?```'
    matches = re.findall(pattern, text, re.DOTALL)

    for match in matches:
        try:
            data = json.loads(match.strip())
            if "name" in data and "arguments" in data:
                calls.append(ToolCall(name=data["name"], arguments=data["arguments"]))
        except json.JSONDecodeError:
            continue

    # Also try inline JSON patterns
    if not calls:
        pattern = r'\{"name":\s*"spatial_query"[^}]+\}'
        matches = re.findall(pattern, text)
        for match in matches:
            try:
                data = json.loads(match)
                if "arguments" in data:
                    calls.append(ToolCall(name=data["name"], arguments=data["arguments"]))
            except json.JSONDecodeError:
                continue

    return calls


def execute_spatial_query(endpoint: str, scene_id: str, query: Dict[str, Any]) -> Dict[str, Any]:
    """Execute a spatial query against the endpoint."""
    # Map tool call arguments to API schema
    api_query = {
        "query_type": query.get("query_type", "distance"),
        "a": query.get("a", "obj_0000"),
        "b": query.get("b"),
        "c": query.get("c"),
        "k": query.get("k", 5),
        "threshold": query.get("threshold", 5.0),
        "tolerance": query.get("tolerance", 2.0),
    }

    payload = {
        "scene_id": scene_id,
        "queries": [api_query],
        "create_scene_if_missing": True,
        "scene_config": {"n_objects": 100, "seed": "gate3_test"}
    }

    try:
        resp = requests.post(f"{endpoint}/echo/spatial", json=payload, timeout=30)
        resp.raise_for_status()
        data = resp.json()
        if data["results"]:
            return data["results"][0]
        return {"error": "No results"}
    except Exception as e:
        return {"error": str(e)}


def run_llm_with_tools(
    llm,
    prompt: str,
    endpoint: str,
    scene_id: str,
    max_tool_rounds: int = 3,
) -> Gate3Result:
    """Run LLM with tool-use loop."""

    full_prompt = f"{SYSTEM_PROMPT}\n\nUser: {prompt}\n\nAssistant:"
    tool_calls = []
    tool_results = []

    for round_num in range(max_tool_rounds):
        # Generate LLM response
        output = llm(
            full_prompt,
            max_tokens=512,
            temperature=0.0,
            top_k=1,
            stop=["User:", "\n\nUser"],
        )

        llm_text = output["choices"][0]["text"].strip()

        # Parse tool calls
        calls = parse_tool_calls(llm_text)

        if not calls:
            # No tool calls - LLM is done
            return Gate3Result(
                prompt=prompt,
                llm_output=llm_text,
                tool_calls=tool_calls,
                tool_results=tool_results,
                final_answer=llm_text,
                passed=True if tool_calls else False,  # Must have used tool at least once
                failure_reason=None if tool_calls else "No tool calls made"
            )

        # Execute tool calls
        for call in calls:
            tool_calls.append(call)
            if call.name == "spatial_query":
                result = execute_spatial_query(endpoint, scene_id, call.arguments)
                tool_results.append(result)

                # Append result to prompt for next round
                result_str = json.dumps(result, indent=2)
                full_prompt += f"\n{llm_text}\n\nTool Result:\n```\n{result_str}\n```\n\nAssistant:"
            else:
                tool_results.append({"error": f"Unknown tool: {call.name}"})

    # Max rounds reached
    return Gate3Result(
        prompt=prompt,
        llm_output=llm_text,
        tool_calls=tool_calls,
        tool_results=tool_results,
        final_answer=llm_text,
        passed=len(tool_calls) > 0,
        failure_reason="Max tool rounds reached" if len(tool_calls) == 0 else None
    )


# Test prompts
TEST_PROMPTS = [
    "Find the 5 nearest objects to obj_0007.",
    "What is the distance between obj_0010 and obj_0050?",
    "Is obj_0020 near obj_0025? Use a threshold of 5.0 units.",
    "Is obj_0005 to the left of obj_0015?",
    "Is obj_0030 above obj_0040?",
    "Find the 3 closest neighbors to obj_0000.",
    "How far is obj_0001 from obj_0099?",
    "Check if obj_0010 is between obj_0000 and obj_0050.",
    "Which objects are closest to obj_0042? Give me 4 of them.",
    "Is the distance from obj_0003 to obj_0007 less than 10 units?",
]


@pytest.fixture(scope="module")
def llm():
    """Load the LLM model once for all tests."""
    if not MODEL_PATH:
        pytest.skip("GATE3_LLM_MODEL not set")

    if not Path(MODEL_PATH).exists():
        pytest.skip(f"Model not found: {MODEL_PATH}")

    try:
        from llama_cpp import Llama
    except ImportError:
        pytest.skip("llama-cpp-python not installed")

    print(f"\nLoading model: {MODEL_PATH}")
    load_start = time.perf_counter()

    model = Llama(
        model_path=MODEL_PATH,
        n_ctx=2048,
        n_threads=8,
        n_gpu_layers=-1,
        seed=42,
        verbose=False,
    )

    load_time = time.perf_counter() - load_start
    print(f"Model loaded in {load_time:.1f}s")

    return model


@pytest.fixture(scope="module")
def endpoint_check():
    """Check if endpoint is healthy."""
    try:
        resp = requests.get(f"{ENDPOINT}/health", timeout=5)
        if resp.status_code != 200:
            pytest.skip(f"Endpoint not healthy: {ENDPOINT}")
    except Exception as e:
        pytest.skip(f"Cannot reach endpoint: {e}")


class TestLLMToolUse:
    """Tests for LLM tool-use integration."""

    @pytest.mark.parametrize("prompt", TEST_PROMPTS[:5])
    def test_tool_call_parsing(self, llm, endpoint_check, prompt):
        """Test that LLM produces parseable tool calls."""
        result = run_llm_with_tools(llm, prompt, ENDPOINT, "llm_test_scene")

        assert result.passed, f"Test failed: {result.failure_reason}"
        assert len(result.tool_calls) > 0, "No tool calls made"

        # Verify tool calls have required fields
        for tc in result.tool_calls:
            assert tc.name == "spatial_query"
            assert "query_type" in tc.arguments
            assert "a" in tc.arguments

    @pytest.mark.parametrize("prompt", TEST_PROMPTS[5:])
    def test_tool_results_used(self, llm, endpoint_check, prompt):
        """Test that LLM uses tool results in final answer."""
        result = run_llm_with_tools(llm, prompt, ENDPOINT, "llm_test_scene")

        assert result.passed, f"Test failed: {result.failure_reason}"

        # Verify tool results were returned
        for tr in result.tool_results:
            assert "error" not in tr or tr.get("error") is None, f"Tool error: {tr}"


class TestGate3Summary:
    """Run full Gate 3 test suite and generate summary."""

    def test_full_gate3_suite(self, llm, endpoint_check):
        """Run all Gate 3 tests and check pass rate."""
        results = []

        for prompt in TEST_PROMPTS:
            result = run_llm_with_tools(llm, prompt, ENDPOINT, "gate3_full_scene")
            results.append(result)

        passed = sum(1 for r in results if r.passed)
        total = len(results)

        # Save results
        output_path = Path(__file__).parent.parent / "receipts" / "gate3_spatial_tool_results.json"
        output_path.parent.mkdir(parents=True, exist_ok=True)

        results_data = {
            "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
            "model": MODEL_PATH,
            "endpoint": ENDPOINT,
            "passed": passed,
            "total": total,
            "tests": [
                {
                    "prompt": r.prompt,
                    "tool_calls": [{"name": tc.name, "arguments": tc.arguments} for tc in r.tool_calls],
                    "tool_results": r.tool_results,
                    "passed": r.passed,
                    "failure_reason": r.failure_reason,
                }
                for r in results
            ]
        }

        with open(output_path, "w") as f:
            json.dump(results_data, f, indent=2)

        print(f"\nResults saved to: {output_path}")
        print(f"Passed: {passed}/{total}")

        # Assert 70%+ pass rate
        assert passed >= total * 0.7, f"Gate 3 failed: only {passed}/{total} tests passed"


if __name__ == "__main__":
    # Command line mode - run full test suite
    if not MODEL_PATH:
        print("ERROR: Set GATE3_LLM_MODEL env var to GGUF model path")
        sys.exit(1)

    if not Path(MODEL_PATH).exists():
        print(f"ERROR: Model not found: {MODEL_PATH}")
        sys.exit(1)

    pytest.main([__file__, "-v", "-s"])
