"""Demo agent using OpenAI SDK — tests Litmus proxy with OpenAI routing."""

from __future__ import annotations

import argparse
import json
import sys


TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "calculate",
            "description": "Perform a math calculation.",
            "parameters": {
                "type": "object",
                "properties": {
                    "expression": {"type": "string", "description": "Math expression"},
                },
                "required": ["expression"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_weather",
            "description": "Get current weather for a city.",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {"type": "string", "description": "City name"},
                },
                "required": ["city"],
            },
        },
    },
]


def handle_tool_call(name: str, args: dict) -> str:
    if name == "calculate":
        try:
            result = eval(args["expression"], {"__builtins__": {}})
            return json.dumps({"result": result})
        except Exception as e:
            return json.dumps({"error": str(e)})
    elif name == "get_weather":
        return json.dumps({
            "city": args.get("city", "Unknown"),
            "temperature": 68,
            "condition": "partly cloudy",
        })
    return json.dumps({"error": f"Unknown tool: {name}"})


def run_agent(prompt: str, stream: bool = False):
    try:
        import openai
    except ImportError:
        print("Error: openai SDK not installed. Run: pip install openai")
        sys.exit(1)

    client = openai.OpenAI()

    print(f"Agent prompt: {prompt}")
    print(f"Streaming: {stream}")
    print("-" * 60)

    messages = [{"role": "user", "content": prompt}]

    for i in range(5):
        print(f"\n[Step {i + 1}] Calling GPT...")

        if stream:
            # Streaming path
            collected_content = ""
            collected_tool_calls = []
            current_tool_call = None

            response_stream = client.chat.completions.create(
                model="gpt-4o-mini",
                max_tokens=256,
                tools=TOOLS,
                messages=messages,
                stream=True,
            )

            finish_reason = None
            for chunk in response_stream:
                delta = chunk.choices[0].delta if chunk.choices else None
                if delta is None:
                    continue

                # Collect text content
                if delta.content:
                    collected_content += delta.content
                    print(delta.content, end="", flush=True)

                # Collect tool calls
                if delta.tool_calls:
                    for tc in delta.tool_calls:
                        if tc.index is not None:
                            while len(collected_tool_calls) <= tc.index:
                                collected_tool_calls.append({
                                    "id": "", "name": "", "arguments": ""
                                })
                            entry = collected_tool_calls[tc.index]
                            if tc.id:
                                entry["id"] = tc.id
                            if tc.function:
                                if tc.function.name:
                                    entry["name"] = tc.function.name
                                if tc.function.arguments:
                                    entry["arguments"] += tc.function.arguments

                if chunk.choices and chunk.choices[0].finish_reason:
                    finish_reason = chunk.choices[0].finish_reason

            if collected_content:
                print()  # newline after streamed text

            print(f"  Finish reason: {finish_reason}")

            if finish_reason == "tool_calls" and collected_tool_calls:
                # Process tool calls
                assistant_msg = {"role": "assistant", "content": collected_content or None, "tool_calls": []}
                tool_results = []
                for tc in collected_tool_calls:
                    args = json.loads(tc["arguments"]) if tc["arguments"] else {}
                    print(f"  Tool call: {tc['name']}({json.dumps(args)})")
                    result = handle_tool_call(tc["name"], args)
                    print(f"  Tool result: {result}")
                    assistant_msg["tool_calls"].append({
                        "id": tc["id"],
                        "type": "function",
                        "function": {"name": tc["name"], "arguments": tc["arguments"]},
                    })
                    tool_results.append({
                        "role": "tool",
                        "tool_call_id": tc["id"],
                        "content": result,
                    })
                messages.append(assistant_msg)
                messages.extend(tool_results)
            else:
                if collected_content:
                    print(f"\n[Final answer]\n{collected_content}")
                break
        else:
            # Non-streaming path
            response = client.chat.completions.create(
                model="gpt-4o-mini",
                max_tokens=256,
                tools=TOOLS,
                messages=messages,
            )

            choice = response.choices[0]
            print(f"  Finish reason: {choice.finish_reason}")

            if choice.finish_reason == "tool_calls" and choice.message.tool_calls:
                tool_results = []
                for tc in choice.message.tool_calls:
                    args = json.loads(tc.function.arguments)
                    print(f"  Tool call: {tc.function.name}({json.dumps(args)})")
                    result = handle_tool_call(tc.function.name, args)
                    print(f"  Tool result: {result}")
                    tool_results.append({
                        "role": "tool",
                        "tool_call_id": tc.id,
                        "content": result,
                    })
                messages.append({"role": "assistant", "content": choice.message.content, "tool_calls": [
                    {"id": tc.id, "type": "function", "function": {"name": tc.function.name, "arguments": tc.function.arguments}}
                    for tc in choice.message.tool_calls
                ]})
                messages.extend(tool_results)
            else:
                print(f"\n[Final answer]\n{choice.message.content}")
                break

    print("\n" + "-" * 60)
    print("Agent finished.")


def main():
    parser = argparse.ArgumentParser(description="Litmus OpenAI demo agent")
    parser.add_argument("--prompt", default="What is 25 * 4?")
    parser.add_argument("--stream", action="store_true", help="Use streaming")
    args = parser.parse_args()
    run_agent(args.prompt, stream=args.stream)


if __name__ == "__main__":
    main()
