#!/usr/bin/env python3
"""
analyze_patient_json_ai_single.py - Single Test AI Analysis (for PHP backend)
"""

import json
import sys
import argparse
import os
from anthropic import Anthropic


ANALYSIS_PROMPT = """Analyze this balance test and return a JSON response with the following structure:

{{
  "summary": "2-3 sentence clinical summary of findings",
  "risk_level": "low|moderate|high",
  "trend": "improving|stable|declining|insufficient_data",
  "icd10_codes": [
    {{"code": "R26.81", "description": "Unsteadiness on feet"}},
    {{"code": "Z91.81", "description": "History of falling"}}
  ],
  "interventions": [
    "Home Safety Checklist review",
    "Physical Therapy referral for balance training",
    "Medication review for fall-risk medications"
  ]
}}

Patient: {name}, Age: {age}
Test Date: {test_date}
Overall Score: {score}/10
Complaints: {complaints}

IMPORTANT: This test ONLY includes the following conditions: {available_conditions}
Only analyze conditions listed above. Do NOT mention or reference any conditions that were not performed.

Test condition codes:
- LBEO = Left leg, Eyes Open
- RBEO = Right leg, Eyes Open
- LBEC = Left leg, Eyes Closed
- RBEC = Right leg, Eyes Closed

Current Test Results:
{results_json}

Previous Tests (for trend analysis):
{previous_tests_json}

IMPORTANT:
- Return ONLY valid JSON, no markdown or explanation
- ICD-10 codes should be relevant to balance/fall risk findings
- Interventions should be specific and actionable
- Risk level based on: low (score 8-10), moderate (score 5-7), high (score 0-4)
- Trend based on comparison with previous tests (use "insufficient_data" if no history)
- This is a MONTHLY testing program. Never recommend "routine testing every 6-12 months" or similar long intervals. Patients are already tested monthly. For follow-up recommendations, use "Continue monthly balance assessments" for stable/good scores, or "Follow-up in 2-4 weeks" only for critical scores below 4.
"""


def analyze_with_ai(test_data, api_key):
    """Use Claude API to analyze test and return structured JSON."""
    try:
        client = Anthropic(api_key=api_key)

        # Parse results_json to find available conditions
        results = json.loads(test_data['results_json'])
        available = []
        condition_names = {
            'lbeo': 'LBEO (Left Eyes Open)',
            'rbeo': 'RBEO (Right Eyes Open)',
            'lbec': 'LBEC (Left Eyes Closed)',
            'rbec': 'RBEC (Right Eyes Closed)'
        }
        for key, name in condition_names.items():
            if results.get(key) is not None:
                available.append(name)

        available_str = ', '.join(available) if available else 'None'

        # Format previous tests for trend analysis
        previous_tests = test_data.get('previous_tests', [])
        if previous_tests:
            previous_tests_json = json.dumps(previous_tests, indent=2)
        else:
            previous_tests_json = "No previous tests available"

        # Format prompt
        prompt = ANALYSIS_PROMPT.format(
            name=test_data['patient_name'],
            age=test_data['age'],
            test_date=test_data.get('test_date', 'Unknown'),
            score=test_data['score'],
            complaints=test_data['complaints'],
            available_conditions=available_str,
            results_json=json.dumps(results, indent=2),
            previous_tests_json=previous_tests_json
        )

        # Call Claude API
        message = client.messages.create(
            model="claude-sonnet-4-20250514",
            max_tokens=1500,
            messages=[
                {"role": "user", "content": prompt}
            ]
        )

        response_text = message.content[0].text

        # Validate it's proper JSON
        try:
            parsed = json.loads(response_text)
            return {
                'success': True,
                'analysis': parsed,  # Return as object, not string
                'model': 'claude-sonnet-4-20250514',
                'tokens': message.usage.input_tokens + message.usage.output_tokens if hasattr(message, 'usage') else None
            }
        except json.JSONDecodeError:
            # Fallback if Claude didn't return valid JSON
            return {
                'success': True,
                'analysis': {
                    'summary': response_text,
                    'risk_level': 'unknown',
                    'trend': 'unknown',
                    'icd10_codes': [],
                    'interventions': []
                },
                'model': 'claude-sonnet-4-20250514',
                'tokens': message.usage.input_tokens + message.usage.output_tokens if hasattr(message, 'usage') else None
            }

    except Exception as e:
        return {
            'success': False,
            'error': f'AI analysis failed: {str(e)}'
        }


def main():
    parser = argparse.ArgumentParser(description='AI-powered patient balance analysis (single test)')
    parser.add_argument('--input', required=True, help='Path to input JSON file')
    args = parser.parse_args()

    try:
        # Get API key
        api_key = os.environ.get('ANTHROPIC_API_KEY')
        if not api_key:
            error_result = {
                'success': False,
                'error': 'ANTHROPIC_API_KEY environment variable not set'
            }
            print(json.dumps(error_result))
            sys.exit(1)

        # Read input
        with open(args.input, 'r') as f:
            test_data = json.load(f)

        # Analyze with AI
        result = analyze_with_ai(test_data, api_key)

        # Output JSON to stdout
        print(json.dumps(result))

    except Exception as e:
        error_result = {
            'success': False,
            'error': f'Failed to process test: {str(e)}'
        }
        print(json.dumps(error_result))
        sys.exit(1)


if __name__ == '__main__':
    main()