#!/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 with your API key
client = anthropic.Anthropic(
    # This is the default and can be omitted
    api_key=os.environ.get("ANTHROPIC_API_KEY"),
)

# Define the prompt for generating the module
def generate_module_prompt(disease_name):
    return f"""
You are a medical expert creating a disease module for the Synthea patient simulation system. 
I need you to create a structured JSON module for {disease_name}.

The module will be used to simulate patients with this condition in the Synthea healthcare simulation system. 
The JSON should follow the Synthea module format which includes:

1. Basic module information (name, remarks)
2. States representing the progression of the disease
3. Transitions between states
4. Guard conditions based on patient attributes where appropriate
5. Care plans and medications that would be prescribed

Your output should be valid JSON that follows the Synthea module structure precisely. Format it as a complete,
well-structured Synthea module. Make sure it's medically accurate and includes all relevant clinical details,
treatment options, and disease progression patterns.

Output only the JSON with no additional commentary or markdown formatting. The output will be directly saved as a file.
"""

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": generate_module_prompt(DISEASE_NAME)}
        ]
    )

    # Extract the JSON from the response
    module_json = message.content[0].text
    
    # 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]
    
    # Parse and format the JSON
    parsed = json.loads(module_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 Exception as e:
    print(f"Error: {e}")
    # Save the raw response for debugging
    with open("error_output.txt", 'w') as f:
        f.write(str(e))
    print("Error details saved to error_output.txt")
    sys.exit(1)