from pathlib import Path
import sys

ROOT = (
    Path.cwd().resolve().parent
    if Path.cwd().name == "notebooks"
    else Path.cwd().resolve()
)
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

import json
from shipit_agent.llms import (
    text_document,
    pdf_document,
    url_pdf_document,
    content_document,
    extract_citations,
)

doc = text_document(
    "The SHIPIT Batch API bills at roughly 50% of standard per-token pricing.",
    title="Pricing FAQ",
    context="Internal pricing doc",
)
print(json.dumps(doc, indent=2))

# Citations can be turned off per document, and PDF/URL variants have the same shape.
print("pdf (citations off):", json.dumps(pdf_document("BASE64DATA", citations=False)))
print("url pdf:           ", json.dumps(url_pdf_document("https://example.com/report.pdf")))
print("content document:  ", json.dumps(content_document("inline grounding text")))

from shipit_agent.llms import AnthropicChatLLM
from shipit_agent.models import Message

llm = AnthropicChatLLM(model="claude-opus-4-1", api_key="offline-demo-key")

req = llm._build_request_kwargs(
    messages=[Message(role="user", content="What discount does the Batch API give?")],
    tools=None,
    system_prompt=None,
    documents=[text_document(
        "The Batch API processes requests asynchronously at ~50% of standard price.",
        title="Pricing FAQ",
    )],
)

last_user = req["messages"][-1]
print("last user message content blocks (document prepended before the question):")
print(json.dumps(last_user["content"], indent=2))

from types import SimpleNamespace as NS

# A fake response: one cited text block + one uncited block.
fake_blocks = [
    NS(type="text", text="The Batch API gives about a 50% discount", citations=[
        {"type": "char_location", "document_index": 0, "document_title": "Pricing FAQ",
         "start_char_index": 0, "end_char_index": 62,
         "cited_text": "processes requests asynchronously at ~50% of standard price"},
    ]),
    NS(type="text", text=" for latency-tolerant work.", citations=None),
]

cites = extract_citations(fake_blocks)
print("metadata['citations'] would be:")
print(json.dumps(cites, indent=2))
print("\nnon-citation responses simply yield:", extract_citations([NS(type="text", text="no cites")]))

from shipit_agent.batch import BatchRequest, BatchResult, BatchRuntime

requests = [
    BatchRequest(custom_id="t1", prompt="Summarise: customer wants a refund.", system="Be terse."),
    BatchRequest(custom_id="t2", prompt="Classify sentiment: I love this product!"),
    BatchRequest(
        custom_id="t3",
        messages=[{"role": "user", "content": "Translate to French: good morning"}],
        model="claude-3-5-haiku-latest",
        max_tokens=128,
    ),
]

print(json.dumps(requests[0].to_payload(), indent=2))

from types import SimpleNamespace as NS

def _text_block(text):
    return NS(type="text", text=text)

def _succeeded(custom_id, text):
    message = NS(
        content=[_text_block(text)],
        usage=NS(input_tokens=12, output_tokens=8),
        stop_reason="end_turn",
    )
    return NS(custom_id=custom_id, result=NS(type="succeeded", message=message))

def _errored(custom_id, message):
    error = NS(error=NS(message=message))
    return NS(custom_id=custom_id, result=NS(type="errored", error=error))

class _FakeBatch:
    def __init__(self, batch_id, status):
        self.id = batch_id
        self.processing_status = status

class _FakeBatches:
    def __init__(self, statuses, results, batch_id="batch_demo"):
        self._statuses = list(statuses)
        self._results = results
        self._batch_id = batch_id
        self.created_with = None
        self.retrieve_calls = 0
    def create(self, *, requests):
        self.created_with = requests
        return _FakeBatch(self._batch_id, self._statuses[0])
    def retrieve(self, batch_id):
        self.retrieve_calls += 1
        idx = min(self.retrieve_calls, len(self._statuses) - 1)
        return _FakeBatch(batch_id, self._statuses[idx])
    def results(self, batch_id):
        return list(self._results)

class _FakeClient:
    def __init__(self, batches):
        self.messages = NS(batches=batches)

batches = _FakeBatches(
    statuses=["in_progress", "in_progress", "ended"],
    results=[
        _succeeded("t1", "Customer requests a refund."),
        _succeeded("t2", "positive"),
        _errored("t3", "rate limited"),
    ],
)
runtime = BatchRuntime(_FakeClient(batches))

sleeps = []
results = runtime.run(requests, poll_interval=30.0, sleep=sleeps.append)  # no-op sleep

print(f"polled {batches.retrieve_calls} times; recorded sleeps: {sleeps}")
for r in results:
    status = "OK" if r.ok else f"ERROR: {r.error}"
    print(f"  {r.custom_id}: {status:<22} output={r.output!r}")

batch_id = runtime.submit(requests)
print("submitted batch id:", batch_id)
print("payload sent to create():")
print(json.dumps(batches.created_with[0], indent=2))

import os
if os.getenv("ANTHROPIC_API_KEY"):
    live = BatchRuntime(api_key=os.environ["ANTHROPIC_API_KEY"])
    bid = live.submit([
        BatchRequest(custom_id="demo", prompt="Say hello in one word.", max_tokens=16),
    ])
    print("submitted real batch:", bid, "status:", live.status(bid))
    print("(batches finish asynchronously; fetch with runtime.results(bid) once ended)")
else:
    print("Set ANTHROPIC_API_KEY to submit a real batch. Skipping (offline).")