#!/usr/bin/env python3

import os
import sys
import json
import anthropic

DISEASE_NAME = "Excessive frequent and irregular menstruation"
OUTPUT_FILE = "excessive_frequent_and_irregular_menstruation.json"

# Initialize the Anthropic client
client = anthropic.Anthropic()

print(f"Generating module for {DISEASE_NAME}...")

try:
    # Send a message to Claude
    message = client.messages.create(
        model="claude-3-7-sonnet-20250219",
        max_tokens=4000,
        temperature=0,
        messages=[
            {"role": "user", "content": f"""Create a Synthea disease module for {DISEASE_NAME} in JSON format. 

The module should follow this structure:
{{
  "name": "Module Name",
  "remarks": [
    "Description of the module"
  ],
  "states": {{
    "Initial": {{
      "type": "Initial",
      "direct_transition": "Next State"
    }},
    "Terminal": {{
      "type": "Terminal"
    }}
    // Additional states with appropriate transitions
  }}
}}

Make sure the JSON is properly formatted with no syntax errors. 
Do not include any markdown formatting, comments, or explanations outside the JSON.
Output only the valid JSON object."""}
        ]
    )

    # Extract the JSON from the response
    module_json = message.content[0].text
    
    # Save the raw response for debugging
    with open(f"{OUTPUT_FILE}.raw", "w") as f:
        f.write(module_json)
        
    print(f"Raw response saved to {OUTPUT_FILE}.raw")
    
    # Find the first { and last } to extract just the JSON part
    start = module_json.find("{")
    end = module_json.rfind("}") + 1
    if start >= 0 and end > start:
        module_json = module_json[start:end]
        
    # Fix common JSON issues
    try:
        # Manual cleaning of known JSON issues
        # Find and remove lines with invalid syntax
        cleaned_lines = []
        for line in module_json.split('\n'):
            # Skip lines with "{%" or any other invalid JSON syntax
            if "{%" in line or "%}" in line or "//" in line:
                print(f"Removing invalid line: {line}")
                continue
            cleaned_lines.append(line)
        
        cleaned_json = '\n'.join(cleaned_lines)
        
        # Try to parse and fix the JSON
        parsed = json.loads(cleaned_json)
        formatted_json = json.dumps(parsed, indent=2)
        
        # Write to file
        with open(OUTPUT_FILE, "w") as f:
            f.write(formatted_json)
        
        print(f"Successfully generated module and saved to {OUTPUT_FILE}")
    except json.JSONDecodeError as e:
        print(f"JSON parsing error: {e}")
        print("Attempting secondary cleaning method...")
        
        # Write the error details for debugging
        with open(f"{OUTPUT_FILE}.error", "w") as f:
            f.write(f"Error: {str(e)}\n\n")
            f.write("JSON that failed to parse:\n")
            f.write(module_json)
    
except Exception as e:
    print(f"Error: {e}")
    sys.exit(1)