"""
Management command to test Number analysis using the AI service.
"""

import json
import logging
from django.core.management.base import BaseCommand
from django.utils import timezone
from bazi.models import Person
from ai.utils.number_analysis import prepare_number_prompt, analyze_number
from ai.services.factory import LLMServiceFactory

logger = logging.getLogger(__name__)

class Command(BaseCommand):
    help = 'Test Number Power AI analysis for a given Person ID'

    def add_arguments(self, parser):
        parser.add_argument('person_id', type=int, help='Person ID to analyze')
        parser.add_argument('--provider', type=str, default='groq', help='LLM provider (groq, openai)')
        parser.add_argument('--model', type=str, default=None, help='Model to use (provider-specific)')
        parser.add_argument('--prompt-only', action='store_true', help='Only output the prompt, do not call LLM')

    def handle(self, *args, **kwargs):
        person_id = kwargs['person_id']
        provider = kwargs.get('provider', 'groq')
        model = kwargs.get('model')
        prompt_only = kwargs.get('prompt_only', False)
        
        try:
            person = Person.objects.get(id=person_id)
        except Person.DoesNotExist:
            self.stderr.write(self.style.ERROR(f'Person with ID {person_id} does not exist'))
            return
        
        if not person.number_result:
            self.stderr.write(self.style.ERROR(f'Person {person.id} has no Number Power result'))
            return
            
        # Prepare prompt
        self.stdout.write('Preparing Number Power analysis prompt...')
        prompt = prepare_number_prompt(person)
        
        if not prompt:
            self.stderr.write(self.style.ERROR('Failed to prepare prompt'))
            return
            
        # Output prompt if requested
        if prompt_only:
            self.stdout.write(self.style.SUCCESS('Number Power Analysis Prompt:'))
            self.stdout.write('-' * 80)
            self.stdout.write(prompt)
            self.stdout.write('-' * 80)
            return
            
        # Use new analyze_number utility to perform and save analysis
        self.stdout.write(f'Using {provider} provider')
        start_time = timezone.now()
        try:
            analysis = analyze_number(person, model_key=model, provider=provider)
        except Exception as e:
            self.stderr.write(self.style.ERROR(f'Error during analysis: {e}'))
            return
        end_time = timezone.now()
        
        if not analysis:
            self.stderr.write(self.style.ERROR('No analysis result'))
            return
        
        # Output results
        self.stdout.write(self.style.SUCCESS('Number Power Analysis Response:'))
        self.stdout.write('-' * 80)
        self.stdout.write(str(analysis.get('number_analysis')))
        self.stdout.write('-' * 80)
        
        # Try to extract JSON from response (if not already parsed)
        data = analysis.get('number_analysis')
        if isinstance(data, str):
            try:
                # Look for JSON block
                json_text = data
                if '```json' in data and '```' in data.split('```json', 1)[1]:
                    json_text = data.split('```json', 1)[1].split('```', 1)[0]
                elif '```' in data and '```' in data.split('```', 1)[1]:
                    json_text = data.split('```', 1)[1].split('```', 1)[0]
                data = json.loads(json_text)
                self.stdout.write(self.style.SUCCESS('Successfully parsed JSON response'))
            except json.JSONDecodeError:
                self.stderr.write(self.style.WARNING('Could not parse JSON from response'))
            except Exception as e:
                self.stderr.write(self.style.WARNING(f'Error processing response: {e}'))
        if isinstance(data, dict):
            sections_count = 0
            for key, section in data.items():
                if isinstance(section, dict) and 'title' in section and 'content' in section:
                    sections_count += 1
                    self.stdout.write(f"- {section.get('title')}")
            self.stdout.write(f"Analysis contains {sections_count} sections")
        duration = (end_time - start_time).total_seconds()
        self.stdout.write(f'Analysis completed in {duration:.2f} seconds')