import os
import sys
from pathlib import Path

from llama_cloud import LlamaCloud
from llama_index.core import Document, VectorStoreIndex
from llama_index.core.embeddings import MockEmbedding


def main() -> None:
    if len(sys.argv) != 2:
        raise SystemExit("Usage: python3 parse_with_llamaparse.py <document-path>")

    if not os.environ.get("LLAMA_CLOUD_API_KEY"):
        raise SystemExit("Set LLAMA_CLOUD_API_KEY before running the parser.")

    source = Path(sys.argv[1])
    if not source.is_file():
        raise SystemExit(f"Input file not found: {source}")

    client = LlamaCloud()
    uploaded = client.files.create(file=source, purpose="parse")
    result = client.parsing.parse(
        file_id=uploaded.id,
        tier="agentic",
        version="latest",
        expand=["markdown_full"],
    )

    markdown = (result.markdown_full or "").strip()
    if not markdown:
        raise SystemExit("Parse completed without markdown output.")

    documents = [
        Document(
            text=markdown,
            metadata={
                "source": source.name,
                "parse_job_id": result.job.id,
                "parse_tier": "agentic",
            },
        )
    ]

    parsed_markdown = Path("parsed-policy-handbook.md")
    parsed_markdown.write_text(markdown + "\n", encoding="utf-8")

    index = VectorStoreIndex.from_documents(
        documents,
        embed_model=MockEmbedding(embed_dim=8),
    )
    retrieved = index.as_retriever(similarity_top_k=1).retrieve("expense approvals")
    preview = markdown.splitlines()[0][:80]

    print(f"parse job: {result.job.status}")
    print(f"llamaindex documents: {len(documents)}")
    print(f"characters: {len(documents[0].text)}")
    print(f"preview: {preview}")
    print(f"indexed nodes: {len(index.index_struct.nodes_dict)}")
    print(f"retrieved nodes: {len(retrieved)}")
    print(f"saved markdown: {parsed_markdown}")


if __name__ == "__main__":
    main()