"""
Lab 1.7: Your First Agent -- SOLUTION

Builds an agent that reads files and answers questions about them.
Uses the Anthropic API (Claude) as the LLM backend, or mock for offline.
"""

import os
import sys

# Import mock LLM for offline execution
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '_shared'))
try:
    from mock_llm import MockAnthropic as Anthropic
    client = Anthropic()
    IS_MOCK = True
except ImportError:
    from anthropic import Anthropic
    client = Anthropic(api_key=os.environ.get("ANTHROPIC_API_KEY"))
    IS_MOCK = False

MAX_ITERATIONS = 10

FILE_READ_TOOL = {
    "name": "read_file",
    "description": "Read the contents of a file at the given path",
    "input_schema": {
        "type": "object",
        "properties": {
            "path": {
                "type": "string",
                "description": "The path to the file to read"
            },
            "reasoning": {
                "type": "string",
                "description": "Why are you reading this file?"
            }
        },
        "required": ["path", "reasoning"]
    }
}

TOOLS = [FILE_READ_TOOL]


def execute_tool(tool_name: str, tool_args: dict) -> str:
    if tool_name == "read_file":
        path = tool_args["path"]
        try:
            with open(path, "r") as f:
                content = f.read()
            return f"Contents of {path}:\n{content}"
        except FileNotFoundError:
            return f"Error: File not found at {path}"
        except Exception as e:
            return f"Error reading file: {str(e)}"
    return f"Unknown tool: {tool_name}"


def run_agent(prompt: str, file_path: str) -> str:
    """Run the agent with a user prompt. Works with real API or mock."""
    
    system_prompt = """You are a helpful assistant with access to file reading tools.
Read the file first, then answer the user's question about its contents.
When you have enough information, provide a clear final answer."""
    
    messages = [
        {"role": "user", "content": f"Read the file at {file_path} and answer: {prompt}"}
    ]
    
    for iteration in range(MAX_ITERATIONS):
        response = client.messages.create(
            model="claude-sonnet-4-20260501" if not IS_MOCK else "mock-model",
            max_tokens=1024,
            system=system_prompt,
            messages=messages,
            tools=TOOLS
        )
        
        if response.stop_reason == "tool_use":
            for block in response.content:
                if block.type == "tool_use":
                    result = execute_tool(block.name, block.input)
                    messages.append({"role": "assistant", "content": response.content})
                    messages.append({"role": "user", "content": result})
        elif response.stop_reason == "end_turn":
            text = "".join(
                block.text for block in response.content if block.type == "text"
            )
            return text
        else:
            return f"Unexpected stop_reason: {response.stop_reason}"
    
    return "Max iterations reached without final answer."


if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("Usage: python solution.py <file_path> <question>")
        sys.exit(1)
    
    file_path = sys.argv[1]
    question = sys.argv[2]
    result = run_agent(question, file_path)
    print(f"\nFinal answer:\n{result}")
