"""
AI-powered migration utility using OpenAI to intelligently fill missing password entry fields.

This module provides:
1. Reading all password entries (metadata only, no passwords)
2. AI inference of missing fields (title, url, username, category, etc.)
3. Generation of a review file for user approval
4. Applying approved changes to the database
"""
import json
import time
import os
import tempfile
import asyncio
from typing import List, Dict, Any, Optional
from onepasswd import ltlog

log = ltlog.getLogger('onepasswd.ai_migrate')

try:
    from openai import AsyncOpenAI
    from pydantic import BaseModel, Field
    OPENAI_AVAILABLE = True
except ImportError:
    OPENAI_AVAILABLE = False
    log.warning("OpenAI SDK not installed. Run: pip install openai")


# Pydantic models for structured output
if OPENAI_AVAILABLE:
    class PasswordEntryMetadata(BaseModel):
        """Structured output for a single password entry."""
        entry_index: int = Field(description="Index of the entry being processed")
        keywords: List[str] = Field(description="Original keywords for the entry")
        title: Optional[str] = Field(None, description="User-friendly display name")
        url: Optional[str] = Field(None, description="Website URL (must use https://)")
        username: Optional[str] = Field(None, description="Username or email address")
        category: Optional[str] = Field(None, description="Category like 'Social', 'Email', 'Work', 'Banking'")
        description: Optional[str] = Field(None, description="Brief description of the account")
        tags: List[str] = Field(default_factory=list, description="Relevant tags")
        confidence: str = Field("medium", description="Confidence level: high, medium, or low")
        notes: Optional[str] = Field(None, description="Explanation of inferences made")

    class PasswordEntriesResponse(BaseModel):
        """Structured output for multiple password entries."""
        entries: List[PasswordEntryMetadata] = Field(description="List of processed entries with metadata")


class AIPasswordMigrator:
    """AI-powered password entry migration assistant."""

    def __init__(self, api_key: Optional[str] = None, model: str = "gpt-4o-mini"):
        """
        Initialize AI migrator.

        Args:
            api_key: OpenAI API key (defaults to OPENAI_API_KEY env var)
            model: OpenAI model to use (default: gpt-4o-mini for cost efficiency)
        """
        if not OPENAI_AVAILABLE:
            raise ImportError("OpenAI SDK not installed. Run: pip install openai")

        self.api_key = api_key or os.environ.get('OPENAI_API_KEY')
        if not self.api_key:
            raise ValueError("OpenAI API key required. Set OPENAI_API_KEY env var or pass api_key parameter")

        self.client = AsyncOpenAI(api_key=self.api_key)
        self.model = model
        log.info(f"Initialized AI migrator with model: {model}")

    def extract_entry_metadata(self, entry) -> Dict[str, Any]:
        """
        Extract metadata from entry without including password.

        Args:
            entry: PasswordEntry instance

        Returns:
            Dict with metadata (no password)
        """
        return {
            'keywords': entry.entries,
            'title': entry.title,
            'url': entry.url,
            'username': entry.username,
            'description': entry.description,
            'category': entry.category,
            'tags': entry.tags,
            'created_time': entry.created_time,
            'modified_time': entry.modified_time
        }

    def identify_missing_fields(self, metadata: Dict[str, Any]) -> List[str]:
        """
        Identify which fields are missing or empty.

        Args:
            metadata: Entry metadata dict

        Returns:
            List of missing field names
        """
        missing = []
        optional_fields = ['title', 'url', 'username', 'description', 'category', 'tags']

        for field in optional_fields:
            value = metadata.get(field)
            if not value or (isinstance(value, (list, str)) and not value):
                missing.append(field)

        return missing

    def build_ai_prompt(self, entries: List[Dict[str, Any]], user_context: str = "") -> str:
        """
        Build AI prompt for filling missing fields.

        Args:
            entries: List of entry metadata dicts
            user_context: Additional user-provided context

        Returns:
            Formatted prompt string
        """
        prompt = """You are a password manager migration assistant. Your task is to intelligently fill missing metadata fields for password entries based on keywords and context.

For each entry, analyze the keywords and existing fields to infer:
- title: A clear, user-friendly display name
- url: The website URL (use https://)
- username: Username or email if inferable from keywords
- category: Category like "Social", "Email", "Work", "Banking", "Shopping", etc.
- description: Brief description of the account
- tags: Relevant tags as a list

Rules:
1. Be conservative - only fill fields you're confident about
2. Use HTTPS URLs (e.g., https://github.com)
3. Capitalize titles properly (e.g., "GitHub Account" not "github account")
4. Leave fields as null if uncertain
5. For known services, use standard URLs (github.com, google.com, etc.)
6. Infer usernames only if keywords contain email or username patterns
7. Categories should be concise (1-2 words)
8. Tags should be relevant and lowercase
9. Set confidence to "high" for well-known services, "medium" for reasonable inferences, "low" for guesses

"""
        if user_context:
            prompt += f"\nUser Context:\n{user_context}\n"

        prompt += f"\nEntries to process ({len(entries)} total):\n\n"

        for idx, entry in enumerate(entries, 1):
            prompt += f"Entry {idx}:\n"
            prompt += f"  Keywords: {', '.join(entry['keywords'])}\n"

            # Show existing fields
            for field in ['title', 'url', 'username', 'category', 'description', 'tags']:
                value = entry.get(field)
                if value and (not isinstance(value, (list, str)) or value):
                    prompt += f"  {field}: {value}\n"

            prompt += "\n"

        prompt += """
Fill the missing metadata for each entry. Respond with structured data following the schema.
"""
        return prompt

    async def call_ai_async(self, prompt: str, entry_index_offset: int = 0) -> List[Dict[str, Any]]:
        """
        Call OpenAI API asynchronously to fill missing fields using structured output.

        Args:
            prompt: Formatted prompt
            entry_index_offset: Offset for entry indices (for batching)

        Returns:
            List of AI-filled entry dicts
        """
        log.info(f"Calling OpenAI API (async, model: {self.model})...")

        try:
            response = await self.client.beta.chat.completions.parse(
                model=self.model,
                messages=[
                    {"role": "system", "content": "You are a helpful password manager migration assistant. Fill missing metadata fields accurately and conservatively."},
                    {"role": "user", "content": prompt}
                ],
                temperature=0.3,  # Lower temperature for more consistent results
                response_format=PasswordEntriesResponse
            )

            result = response.choices[0].message.parsed
            log.info(f"AI response received: {len(result.entries)} entries processed")

            # Convert Pydantic models to dicts
            entries_data = []
            for entry in result.entries:
                entries_data.append({
                    'entry_index': entry.entry_index,
                    'keywords': entry.keywords,
                    'title': entry.title,
                    'url': entry.url,
                    'username': entry.username,
                    'category': entry.category,
                    'description': entry.description,
                    'tags': entry.tags,
                    'confidence': entry.confidence,
                    'notes': entry.notes
                })

            return entries_data

        except Exception as e:
            log.error(f"AI API call failed: {e}")
            raise

    def call_ai(self, prompt: str) -> List[Dict[str, Any]]:
        """
        Synchronous wrapper for call_ai_async.

        Args:
            prompt: Formatted prompt

        Returns:
            List of AI-filled entry dicts
        """
        return asyncio.run(self.call_ai_async(prompt))

    async def process_entries_async(
        self,
        db,
        user_context: str = "",
        batch_size: int = 20,
        max_concurrent: int = 3,
        progress_callback: Optional[callable] = None
    ) -> Dict[str, Any]:
        """
        Process all entries that need migration asynchronously with rate limiting.

        Args:
            db: Database instance
            user_context: User-provided context for better AI inference
            batch_size: Number of entries to process per API call
            max_concurrent: Maximum number of concurrent API calls (default: 3)
            progress_callback: Optional callback(batch_num, total_batches, result) for progress updates

        Returns:
            Dict with original and proposed entries
        """
        from onepasswd.models import PasswordEntry

        # Collect entries that need migration
        entries_to_process = []

        for entry in db.get_all_items():
            metadata = self.extract_entry_metadata(entry)
            missing = self.identify_missing_fields(metadata)

            if missing:
                entries_to_process.append({
                    'original_entry': entry,
                    'metadata': metadata,
                    'missing_fields': missing
                })

        log.info(f"Found {len(entries_to_process)} entries needing migration")

        if not entries_to_process:
            return {'entries': [], 'message': 'No entries need migration'}

        # Create batches
        batches = []
        for i in range(0, len(entries_to_process), batch_size):
            batch = entries_to_process[i:i + batch_size]
            batches.append(batch)

        log.info(f"Processing {len(batches)} batches with max {max_concurrent} concurrent requests")

        # Semaphore to limit concurrent API calls
        semaphore = asyncio.Semaphore(max_concurrent)

        # Process batches with rate limiting
        async def process_batch(batch_num: int, batch: List[Dict]) -> List[Dict]:
            async with semaphore:
                batch_metadata = [item['metadata'] for item in batch]

                log.info(f"[Batch {batch_num + 1}/{len(batches)}] Processing {len(batch)} entries...")

                # Build prompt and call AI
                prompt = self.build_ai_prompt(batch_metadata, user_context)

                # Retry logic for rate limits
                max_retries = 3
                retry_delay = 2.0

                for attempt in range(max_retries):
                    try:
                        ai_results = await self.call_ai_async(prompt, entry_index_offset=batch_num * batch_size)
                        break
                    except Exception as e:
                        if "rate" in str(e).lower() or "429" in str(e):
                            if attempt < max_retries - 1:
                                wait_time = retry_delay * (2 ** attempt)  # Exponential backoff
                                log.warning(f"[Batch {batch_num + 1}] Rate limit hit, retrying in {wait_time}s...")
                                await asyncio.sleep(wait_time)
                            else:
                                log.error(f"[Batch {batch_num + 1}] Rate limit exceeded after {max_retries} attempts")
                                raise
                        else:
                            raise

                # Combine original and AI-filled data
                batch_results = []
                for idx, item in enumerate(batch):
                    ai_result = ai_results[idx] if idx < len(ai_results) else {}

                    result = {
                        'original': item['metadata'],
                        'proposed': ai_result,
                        'missing_fields': item['missing_fields'],
                        'original_entry': item['original_entry']
                    }
                    batch_results.append(result)

                log.info(f"[Batch {batch_num + 1}/{len(batches)}] ✓ Completed")

                # Call progress callback if provided
                if progress_callback:
                    progress_callback(batch_num + 1, len(batches), batch_results)

                return batch_results

        # Run batches with controlled concurrency
        tasks = [process_batch(i, batch) for i, batch in enumerate(batches)]
        all_batch_results = await asyncio.gather(*tasks)

        # Flatten results
        all_results = []
        for batch_results in all_batch_results:
            all_results.extend(batch_results)

        return {
            'entries': all_results,
            'total_processed': len(all_results),
            'message': f'Successfully processed {len(all_results)} entries'
        }

    def process_entries(self, db, user_context: str = "", batch_size: int = 20) -> Dict[str, Any]:
        """
        Synchronous wrapper for process_entries_async.

        Args:
            db: Database instance
            user_context: User-provided context for better AI inference
            batch_size: Number of entries to process per API call

        Returns:
            Dict with original and proposed entries
        """
        return asyncio.run(self.process_entries_async(db, user_context, batch_size))

    def generate_review_file(self, results: Dict[str, Any], output_path: str = None) -> str:
        """
        Generate human-readable review file.

        Args:
            results: Results from process_entries()
            output_path: Output file path (defaults to temp file)

        Returns:
            Path to review file
        """
        if output_path is None:
            fd, output_path = tempfile.mkstemp(suffix='.json', prefix='onepasswd_ai_review_')
            os.close(fd)

        review_data = {
            'generated_at': time.time(),
            'generated_at_human': time.ctime(),
            'total_entries': results.get('total_processed', 0),
            'instructions': {
                'review': 'Review each entry below. Modify proposed fields as needed.',
                'apply': 'Run: onepasswd ai-migrate apply <this_file>',
                'format': 'Edit the "proposed" fields. Set to null to leave unchanged.'
            },
            'entries': []
        }

        for item in results.get('entries', []):
            entry_data = {
                'keywords': item['original']['keywords'],
                'missing_fields': item['missing_fields'],
                'original': {
                    k: v for k, v in item['original'].items()
                    if k != 'keywords' and v
                },
                'proposed': item['proposed'],
                'approved': True,  # User can set to False to skip
                'notes': ''  # User can add notes
            }
            review_data['entries'].append(entry_data)

        # Write to file with pretty formatting
        with open(output_path, 'w') as f:
            json.dump(review_data, f, indent=2, ensure_ascii=False)

        log.info(f"Review file generated: {output_path}")
        return output_path

    def apply_review_file(self, db, review_file_path: str, master_password: str):
        """
        Apply approved changes from review file to database.

        Args:
            db: Database instance
            review_file_path: Path to reviewed JSON file
            master_password: Master password for encryption
        """
        from onepasswd.models import PasswordEntry
        from onepasswd.onepasswd import DB

        # Load review file
        with open(review_file_path, 'r') as f:
            review_data = json.load(f)

        applied_count = 0
        skipped_count = 0
        error_count = 0

        log.info(f"Applying changes from {review_file_path}")

        for entry_data in review_data.get('entries', []):
            # Skip if not approved
            if not entry_data.get('approved', True):
                skipped_count += 1
                continue

            keywords = tuple(entry_data['keywords'])
            proposed = entry_data.get('proposed', {})

            try:
                # Build update dict (only non-null values)
                updates = {}
                for field in ['title', 'url', 'username', 'description', 'category', 'tags']:
                    value = proposed.get(field)
                    if value is not None and value != '':
                        updates[field] = value

                if updates:
                    db.update_item(keywords, **updates)
                    applied_count += 1
                    log.info(f"Updated entry: {keywords}")
                else:
                    skipped_count += 1
                    log.info(f"Skipped entry (no updates): {keywords}")

            except Exception as e:
                error_count += 1
                log.error(f"Failed to update entry {keywords}: {e}")

        log.info(f"Migration complete: {applied_count} applied, {skipped_count} skipped, {error_count} errors")

        return {
            'applied': applied_count,
            'skipped': skipped_count,
            'errors': error_count
        }


def check_openai_available() -> bool:
    """Check if OpenAI SDK is installed."""
    return OPENAI_AVAILABLE
