#!/usr/bin/env python3
"""
Simple Module Generator for Synthea Nextflow Pipeline
Using Direct HTTP Requests to avoid client library issues
"""

import os
import json
import argparse
import requests
from dotenv import load_dotenv

def generate_module(disease_name, output_file):
    """Generate a Synthea module for the specified disease"""
    # Load API key from environment
    load_dotenv()
    api_key = os.getenv("ANTHROPIC_API_KEY")
    
    if not api_key:
        raise ValueError("ANTHROPIC_API_KEY environment variable not set")
    
    print(f"Generating module for {disease_name}...")
    
    # Current Anthropic API endpoint and format (as of 2024)
    url = "https://api.anthropic.com/v1/messages"
    
    headers = {
        "Content-Type": "application/json",
        "x-api-key": api_key,
        "anthropic-version": "2023-06-01"
    }
    
    system_prompt = """
    You are an expert in medical informatics and Synthea module creation. 
    Generate a complete, valid JSON module for the specified disease.
    The module must follow Synthea's format conventions.
    """
    
    user_prompt = f"""
    Create a complete Synthea module for {disease_name}.
    
    The module should include:
    - Initial states for disease onset and progression
    - Diagnostic procedures and criteria
    - Treatment options and medication regimens
    - Complications and their management
    - Follow-up care protocols
    
    Return ONLY valid JSON that can be directly used in Synthea without any explanation or markdown.
    """
    
    data = {
        "model": "claude-3-opus-20240229",
        "system": system_prompt,
        "messages": [
            {
                "role": "user",
                "content": user_prompt
            }
        ],
        "max_tokens": 4000,
        "temperature": 0.2
    }
    
    try:
        # Make direct API request
        response = requests.post(url, headers=headers, json=data)
        
        # Check for errors
        if response.status_code != 200:
            print(f"API request failed with status code {response.status_code}: {response.text}")
            raise Exception(f"API request failed with status code {response.status_code}")
        
        # Parse response
        result = response.json()
        module_content = result["content"][0]["text"]
        
        # Extract the JSON part if wrapped in markdown
        if "```json" in module_content:
            module_content = module_content.split("```json")[1].split("```")[0].strip()
        elif "```" in module_content:
            module_content = module_content.split("```")[1].split("```")[0].strip()
            
        # Validate JSON
        try:
            module_json = json.loads(module_content)
            with open(output_file, 'w') as f:
                json.dump(module_json, f, indent=2)
            print(f"Successfully generated module and saved to {output_file}")
            return True
        except json.JSONDecodeError as e:
            print(f"Generated content is not valid JSON: {e}")
            with open(f"{output_file}.raw", 'w') as f:
                f.write(module_content)
            print(f"Raw content saved to {output_file}.raw")
            return False
            
    except Exception as e:
        print(f"Error generating module: {e}")
        if isinstance(e, requests.exceptions.RequestException):
            print(f"Request error details: {e.response.text if hasattr(e, 'response') else 'No response details'}")
        return False

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Generate Synthea disease module')
    parser.add_argument('--disease', required=True, help='Disease name to generate module for')
    parser.add_argument('--output', required=True, help='Output filename for the module')
    
    args = parser.parse_args()
    
    # Install required packages if needed
    try:
        import requests
    except ImportError:
        import subprocess
        print("Installing required packages...")
        subprocess.check_call(["pip", "install", "requests"])
        
    success = generate_module(args.disease, args.output)
    
    if not success:
        exit(1)