import re

with open("src/orchestrator_pm.py", "r", encoding="utf-8", newline="") as f:
    content = f.read()

# Replace the JSON parsing section
old = """   # 4. Parse JSON Output
   try:
   # The prompt asks for a JSON array. We need to extract it if the AI added markdown blocks.
    json_match = response.strip()
    if "```json" in json_match:
     json_match = json_match.split("```json")[1].split("```")[0].strip()
    elif "```" in json_match:
     json_match = json_match.split("```")[1].split("```")[0].strip()
    tracks: list[dict[str, Any]] = json.loads(json_match)"""

new = """   # 4. Parse JSON Output
   try:
    json_match = response.strip()
    # Handle mock_gemini_cli.py format: {"type": "message", "content": "[...]"}
    if '"content": "' in json_match:
        match = re.search(r'"content"\\s*:\\s*"(\\[.*?\\])"', json_match)
        if match:
            json_match = match.group(1)
    # Handle markdown code blocks
    if "```json" in json_match:
        json_match = json_match.split("```json")[1].split("```")[0].strip()
    elif "```" in json_match:
        json_match = json_match.split("```")[1].split("```")[0].strip()
    tracks: list[dict[str, Any]] = json.loads(json_match)"""

content = content.replace(old, new)

with open("src/orchestrator_pm.py", "w", encoding="utf-8", newline="") as f:
    f.write(content)

print("Fixed JSON parsing in orchestrator_pm.py")