import re
import json
import argparse
from pathlib import Path

def _split_blocks(text):
    parts = re.split(r"^={80}\n", text, flags=re.M)
    return [p.strip() for p in parts if p.strip()]


def _extract_query(text):
    query_match = re.search(r"^Query:\s*(.+)$", text, re.M)
    if not query_match:
        return ""
    return query_match.group(1).strip()


def parse_output_txt(path, query=None):
    with open(path, "r", encoding="utf-8") as f:
        text = f.read()

    blocks = _split_blocks(text)
    if not blocks:
        raise ValueError("No output blocks found in output.txt")

    chosen_block = blocks[-1]
    if query:
        for block in reversed(blocks):
            if _extract_query(block) == query:
                chosen_block = block
                break

    text = chosen_block

    # Extract query
    query_match = re.search(r"^Query:\s*(.+)$", text, re.M)
    query = query_match.group(1).strip()

    # Extract chunks
    chunk_pattern = re.compile(
        r"\[(\d+)\]\s+Score:\s+([\d.]+)\s+\|\s+ID:\s+([^\n]+)\n\s+(.*?)(?=\n\n|\Z)",
        re.S
    )

    results = []
    for _, score, chunk_id, content in chunk_pattern.findall(text):
        results.append({
            "chunk_id": chunk_id.strip(),
            "score": float(score),
            "content": content.strip()
        })

    return {
        "query": query,
        "retrieved_results": results
    }

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Parse the latest output block.")
    parser.add_argument("--query", help="Query text to match in output.txt")
    args = parser.parse_args()

    parsed = parse_output_txt("../04_vectoredb/output.txt", query=args.query)

    with open("parsed_input.json", "w") as f:
        json.dump(parsed, f, indent=2)

    print("Parsed input written to parsed_input.json")