import json
import os
import sys
from typing import Any, Dict, List

import requests

# Required: export DATABRICKS_TOKEN with a valid Databricks PAT before running.
DATABRICKS_TOKEN = "dapied25006cd5ec824b138dfa5e82e2630c"   #User1

if not DATABRICKS_TOKEN:
    raise RuntimeError("Set DATABRICKS_TOKEN with a valid Databricks PAT.")

# Optional: override DATABRICKS_BASE_URL to point at your workspace host.
BASE_URL = os.environ.get(
    "DATABRICKS_BASE_URL",
    "https://adb-7405604858124224.4.azuredatabricks.net",
)
ENDPOINT = f"{BASE_URL.rstrip('/')}/serving-endpoints/anthropic/v1/messages"
# Adjust DEFAULT_MODEL if your workspace exposes a different deployed model name.
DEFAULT_MODEL = "databricks-claude-sonnet-4-5"


def build_payload(prompt: str, model: str, max_tokens: int) -> Dict[str, Any]:
    return {
        "model": model,
        "max_tokens": max_tokens,
        "messages": [{"role": "user", "content": prompt}],
    }


def call_message_api(
    prompt: str, model: str = DEFAULT_MODEL, max_tokens: int = 500
) -> Dict[str, Any]:
    payload = build_payload(prompt, model, max_tokens)
    headers = {
        "Authorization": f"Bearer {DATABRICKS_TOKEN}",
        "Content-Type": "application/json",
    }
    response = requests.post(ENDPOINT, headers=headers, json=payload, timeout=30)
    response.raise_for_status()
    return response.json()


def extract_text(content: Any) -> str:
    if isinstance(content, list):
        text_blocks: List[str] = []
        for block in content:
            if isinstance(block, dict) and block.get("type") == "text":
                text_blocks.append(str(block.get("text", "")))
        if text_blocks:
            return "\n".join(text_blocks)
    return str(content)


def main(argv: List[str]) -> None:
    prompt = " ".join(argv) if argv else "What is an LLM agent?"
    result = call_message_api(prompt)
    content = result.get("content")
    output = extract_text(content)
    print(output or json.dumps(result, indent=2))


if __name__ == "__main__":
    main(sys.argv[1:])