"""Parser for Automatic1111 metadata format."""

import re
import os
import json
import logging
from typing import Dict, Any
from ..base import RecipeMetadataParser
from ..constants import GEN_PARAM_KEYS
from ...services.metadata_service import get_default_metadata_provider

logger = logging.getLogger(__name__)

class AutomaticMetadataParser(RecipeMetadataParser):
    """Parser for Automatic1111 metadata format"""
    
    METADATA_MARKER = r"Steps: \d+"
    
    # Regular expressions for extracting specific metadata
    HASHES_REGEX = r', Hashes:\s*({[^}]+})'
    LORA_HASHES_REGEX = r', Lora hashes:\s*"([^"]+)"'
    CIVITAI_RESOURCES_REGEX = r', Civitai resources:\s*(\[\{.*?\}\])'
    CIVITAI_METADATA_REGEX = r', Civitai metadata:\s*(\{.*?\})'
    EXTRANETS_REGEX = r'<(lora|hypernet):([^:]+):(-?[0-9.]+)>'
    MODEL_HASH_PATTERN = r'Model hash: ([a-zA-Z0-9]+)'
    MODEL_NAME_PATTERN = r'Model: ([^,]+)'
    VAE_HASH_PATTERN = r'VAE hash: ([a-zA-Z0-9]+)'
    
    def is_metadata_matching(self, user_comment: str) -> bool:
        """Check if the user comment matches the Automatic1111 format"""
        return re.search(self.METADATA_MARKER, user_comment) is not None
    
    async def parse_metadata(self, user_comment: str, recipe_scanner=None, civitai_client=None) -> Dict[str, Any]:
        """Parse metadata from Automatic1111 format"""
        try:
            # Get metadata provider instead of using civitai_client directly
            metadata_provider = await get_default_metadata_provider()
            
            # Split on Negative prompt if it exists
            if "Negative prompt:" in user_comment:
                parts = user_comment.split('Negative prompt:', 1)
                prompt = parts[0].strip()
                negative_and_params = parts[1] if len(parts) > 1 else ""
            else:
                # No negative prompt section
                param_start = re.search(self.METADATA_MARKER, user_comment)
                if param_start:
                    prompt = user_comment[:param_start.start()].strip()
                    negative_and_params = user_comment[param_start.start():]
                else:
                    prompt = user_comment.strip()
                    negative_and_params = ""
            
            # Initialize metadata
            metadata: Dict[str, Any] = {
                "prompt": prompt,
                "loras": []
            }
            
            # Extract negative prompt and parameters
            if negative_and_params:
                # If we split on "Negative prompt:", check for params section
                if "Negative prompt:" in user_comment:
                    param_start = re.search(r'Steps: ', negative_and_params)
                    if param_start:
                        neg_prompt = negative_and_params[:param_start.start()].strip()
                        metadata["negative_prompt"] = neg_prompt
                        params_section = negative_and_params[param_start.start():]
                    else:
                        metadata["negative_prompt"] = negative_and_params.strip()
                        params_section = ""
                else:
                    # No negative prompt, entire section is params
                    params_section = negative_and_params
                
                # Extract generation parameters
                if params_section:
                    # Extract Civitai resources
                    civitai_resources_match = re.search(self.CIVITAI_RESOURCES_REGEX, params_section)
                    if civitai_resources_match:
                        try:
                            civitai_resources = json.loads(civitai_resources_match.group(1))
                            metadata["civitai_resources"] = civitai_resources
                            params_section = params_section.replace(civitai_resources_match.group(0), '')
                        except json.JSONDecodeError:
                            logger.error("Error parsing Civitai resources JSON")
                    
                    # Extract Hashes
                    hashes_match = re.search(self.HASHES_REGEX, params_section)
                    if hashes_match:
                        try:
                            hashes = json.loads(hashes_match.group(1))
                            # Process hash keys
                            processed_hashes = {}
                            for key, value in hashes.items():
                                # Convert Model: or LORA: prefix to lowercase if present
                                if ':' in key:
                                    prefix, name = key.split(':', 1)
                                    prefix = prefix.lower()
                                else:
                                    prefix = ''
                                    name = key

                                # Clean up the name part
                                if '/' in name:
                                    name = name.split('/')[-1]  # Get last part after /
                                if '.safetensors' in name:
                                    name = name.split('.safetensors')[0]  # Remove .safetensors
                                
                                # Reconstruct the key
                                new_key = f"{prefix}:{name}" if prefix else name
                                processed_hashes[new_key] = value

                            metadata["hashes"] = processed_hashes
                            # Remove hashes from params section to not interfere with other parsing
                            params_section = params_section.replace(hashes_match.group(0), '')
                        except json.JSONDecodeError:
                            logger.error("Error parsing hashes JSON")
                    
                    # Pick up model hash from parsed hashes if available
                    if "hashes" in metadata and not metadata.get("model_hash"):
                        model_hash_from_hashes = metadata["hashes"].get("model")
                        if model_hash_from_hashes:
                            metadata["model_hash"] = model_hash_from_hashes
                    
                    # Extract Lora hashes in alternative format.
                    # Run unconditionally (not just as fallback) so that
                    # non-empty hashes from Lora hashes fill in the gaps left
                    # by empty values in the Hashes JSON dict.  Some WebUI
                    # builds write real hash values only to Lora hashes and
                    # leave the Hashes JSON values empty.
                    lora_hashes_match = re.search(self.LORA_HASHES_REGEX, params_section)
                    if lora_hashes_match:
                        try:
                            lora_hashes_str = lora_hashes_match.group(1)
                            lora_hash_entries = lora_hashes_str.split(', ')

                            # Parse each lora hash entry (format: "name: hash")
                            for entry in lora_hash_entries:
                                if ': ' in entry:
                                    lora_name, lora_hash = entry.split(': ', 1)
                                    lora_hash = lora_hash.strip()
                                    if not lora_hash:
                                        # Skip entries without a hash value
                                        continue
                                    # Initialize hashes dict if it doesn't exist
                                    if "hashes" not in metadata:
                                        metadata["hashes"] = {}
                                    # Add as lora type in the same format as
                                    # regular hashes.  Only override an
                                    # existing entry if its value is empty
                                    # (Lora hashes is the more reliable
                                    # source when Hashes JSON has blanks).
                                    key = f"lora:{lora_name}"
                                    existing = metadata["hashes"].get(key, "")
                                    if not existing:
                                        metadata["hashes"][key] = lora_hash

                            # Remove lora hashes from params section
                            params_section = params_section.replace(lora_hashes_match.group(0), '')
                        except Exception as e:
                            logger.error(f"Error parsing Lora hashes: {e}")

                    # Extract checkpoint model hash/name when provided outside Civitai resources
                    model_hash_match = re.search(self.MODEL_HASH_PATTERN, params_section)
                    if model_hash_match:
                        metadata["model_hash"] = model_hash_match.group(1).strip()
                        params_section = params_section.replace(model_hash_match.group(0), '')

                    model_name_match = re.search(self.MODEL_NAME_PATTERN, params_section)
                    if model_name_match:
                        metadata["model_name"] = model_name_match.group(1).strip()
                        params_section = params_section.replace(model_name_match.group(0), '')
                    
                    # Extract basic parameters
                    param_pattern = r'([A-Za-z\s]+): ([^,]+)'
                    params = re.findall(param_pattern, params_section)
                    gen_params = {}
                    
                    for key, value in params:
                        clean_key = key.strip().lower().replace(' ', '_')
                        
                        # Skip if not in recognized gen param keys
                        if clean_key not in GEN_PARAM_KEYS:
                            continue
                            
                        # Convert numeric values
                        if clean_key in ['steps', 'seed']:
                            try:
                                gen_params[clean_key] = int(value.strip())
                            except ValueError:
                                gen_params[clean_key] = value.strip()
                        elif clean_key in ['cfg_scale']:
                            try:
                                gen_params[clean_key] = float(value.strip())
                            except ValueError:
                                gen_params[clean_key] = value.strip()
                        else:
                            gen_params[clean_key] = value.strip()
                    
                    # Extract size if available and add to gen_params if a recognized key
                    size_match = re.search(r'Size: (\d+)x(\d+)', params_section)
                    if size_match and 'size' in GEN_PARAM_KEYS:
                        width, height = size_match.groups()
                        gen_params['size'] = f"{width}x{height}"
                    
                    # Add prompt and negative_prompt to gen_params if they're in GEN_PARAM_KEYS
                    if 'prompt' in GEN_PARAM_KEYS and 'prompt' in metadata:
                        gen_params['prompt'] = metadata['prompt']
                    if 'negative_prompt' in GEN_PARAM_KEYS and 'negative_prompt' in metadata:
                        gen_params['negative_prompt'] = metadata['negative_prompt']
                    
                    metadata["gen_params"] = gen_params
            
            # Extract LoRA and checkpoint information 
            loras = []
            base_model_counts = {}
            checkpoint = None
            
            # First use Civitai resources if available (more reliable source)
            if metadata.get("civitai_resources"):
                for resource in metadata.get("civitai_resources", []):
                    # --- Added: Parse 'air' field if present ---
                    air = resource.get("air")
                    if air:
                        # Format: urn:air:sdxl:lora:civitai:1221007@1375651
                        # Or: urn:air:sdxl:checkpoint:civitai:623891@2019115
                        air_pattern = r"urn:air:[^:]+:(?P<type>[^:]+):civitai:(?P<modelId>\d+)@(?P<modelVersionId>\d+)"
                        air_match = re.match(air_pattern, air)
                        if air_match:
                            air_type = air_match.group("type")
                            air_modelId = int(air_match.group("modelId"))
                            air_modelVersionId = int(air_match.group("modelVersionId"))
                            # checkpoint/lycoris/lora/hypernet
                            resource["type"] = air_type
                            resource["modelId"] = air_modelId
                            resource["modelVersionId"] = air_modelVersionId
                    # --- End added ---

                    if resource.get("type") == "checkpoint" and resource.get("modelVersionId"):
                        version_id = resource.get("modelVersionId")
                        version_id_str = str(version_id)
                        checkpoint_entry = {
                            'id': version_id,
                            'modelId': resource.get("modelId", 0),
                            'name': resource.get("modelName", "Unknown Checkpoint"),
                            'version': resource.get("modelVersionName", resource.get("versionName", "")),
                            'type': resource.get("type", "checkpoint"),
                            'existsLocally': False,
                            'localPath': None,
                            'file_name': resource.get("modelName", ""),
                            'hash': resource.get("hash", "") or "",
                            'thumbnailUrl': '/loras_static/images/no-preview.png',
                            'baseModel': '',
                            'size': 0,
                            'downloadUrl': '',
                            'isDeleted': False
                        }

                        if metadata_provider:
                            try:
                                civitai_info = await metadata_provider.get_model_version_info(version_id_str)
                                checkpoint_entry = await self.populate_checkpoint_from_civitai(
                                    checkpoint_entry,
                                    civitai_info
                                )
                            except Exception as e:
                                logger.error(
                                    "Error fetching Civitai info for checkpoint version %s: %s",
                                    version_id,
                                    e,
                                )

                        # Prefer the first checkpoint found
                        if checkpoint_entry.get("baseModel"):
                            base_model_value = checkpoint_entry["baseModel"]
                            base_model_counts[base_model_value] = base_model_counts.get(base_model_value, 0) + 1

                        if checkpoint is None:
                            checkpoint = checkpoint_entry

                        continue

                    if resource.get("type") in ["lora", "lycoris", "hypernet"] and resource.get("modelVersionId"):
                        # Initialize lora entry
                        lora_entry = {
                            'id': resource.get("modelVersionId", 0),
                            'modelId': resource.get("modelId", 0),
                            'name': resource.get("modelName", "Unknown LoRA"),
                            'version': resource.get("modelVersionName", resource.get("versionName", "")),
                            'type': resource.get("type", "lora"),
                            'weight': round(float(resource.get("weight", 1.0)), 2),
                            'existsLocally': False,
                            'thumbnailUrl': '/loras_static/images/no-preview.png',
                            'baseModel': '',
                            'size': 0,
                            'downloadUrl': '',
                            'isDeleted': False
                        }
                        
                        # Get additional info from Civitai
                        if metadata_provider:
                            try:
                                civitai_info = await metadata_provider.get_model_version_info(resource.get("modelVersionId"))
                                populated_entry = await self.populate_lora_from_civitai(
                                    lora_entry,
                                    civitai_info,
                                    recipe_scanner,
                                    base_model_counts
                                )
                                if populated_entry is None:
                                    continue  # Skip invalid LoRA types
                                lora_entry = populated_entry
                            except Exception as e:
                                logger.error(f"Error fetching Civitai info for LoRA {lora_entry['name']}: {e}")
                        
                        loras.append(lora_entry)
            
            # Fallback checkpoint parsing from generic "Model" and "Model hash" fields
            if checkpoint is None:
                model_hash = metadata.get("model_hash")
                if not model_hash and metadata.get("hashes"):
                    model_hash = metadata["hashes"].get("model")

                model_name = metadata.get("model_name")
                file_name = ""
                if model_name:
                    cleaned_name = re.split(r"[\\\\/]", model_name)[-1]
                    file_name = os.path.splitext(cleaned_name)[0]

                if model_hash or model_name:
                    checkpoint_entry = {
                        'id': 0,
                        'modelId': 0,
                        'name': model_name or "Unknown Checkpoint",
                        'version': '',
                        'type': 'checkpoint',
                        'hash': model_hash or "",
                        'existsLocally': False,
                        'localPath': None,
                        'file_name': file_name,
                        'thumbnailUrl': '/loras_static/images/no-preview.png',
                        'baseModel': '',
                        'size': 0,
                        'downloadUrl': '',
                        'isDeleted': False
                    }

                    if metadata_provider and model_hash:
                        try:
                            civitai_info = await metadata_provider.get_model_by_hash(model_hash)
                            checkpoint_entry = await self.populate_checkpoint_from_civitai(
                                checkpoint_entry,
                                civitai_info
                            )
                        except Exception as e:
                            logger.error(f"Error fetching Civitai info for checkpoint hash {model_hash}: {e}")

                    if checkpoint_entry.get("baseModel"):
                        base_model_value = checkpoint_entry["baseModel"]
                        base_model_counts[base_model_value] = base_model_counts.get(base_model_value, 0) + 1

                    checkpoint = checkpoint_entry

            def normalize_lora_name(name, basename=False):
                normalized = str(name or '').replace('\\', '/')
                if normalized.casefold().endswith('.safetensors'):
                    normalized = normalized[:-12]
                if basename:
                    normalized = normalized.rsplit('/', 1)[-1]
                return normalized.casefold()

            def get_version_id(lora):
                version_id = lora.get('id')
                if version_id in (None, '', 0, '0'):
                    version_id = lora.get('modelVersionId')
                if version_id in (None, '', 0, '0'):
                    return None
                return str(version_id)

            prompt_loras = {}
            for match in re.findall(self.EXTRANETS_REGEX, prompt):
                lora_type, lora_name, _ = match
                prompt_loras[(lora_type, normalize_lora_name(lora_name))] = match

            prompt_by_basename = {}
            for lora_type, lora_name, lora_weight in prompt_loras.values():
                key = (lora_type, normalize_lora_name(lora_name, True))
                prompt_by_basename.setdefault(key, []).append((lora_name, round(float(lora_weight), 2)))

            hash_basenames = {
                (hash_key.split(':', 1)[0], normalize_lora_name(hash_key.split(':', 1)[1], True))
                for hash_key, hash_value in metadata.get("hashes", {}).items()
                if hash_value and hash_key.startswith(("lora:", "hypernet:"))
            }
            recipe_base_model = checkpoint.get("baseModel") if checkpoint else None
            if not recipe_base_model and len(base_model_counts) == 1:
                recipe_base_model = next(iter(base_model_counts))

            resource_lora_count = len(loras)

            def make_lora_entry(lora_type, lora_name, weight, lora_hash=''):
                return {
                    'name': lora_name,
                    'type': lora_type,
                    'weight': weight,
                    'hash': lora_hash,
                    'existsLocally': False,
                    'localPath': None,
                    'file_name': lora_name,
                    'thumbnailUrl': '/loras_static/images/no-preview.png',
                    'baseModel': '',
                    'size': 0,
                    'downloadUrl': '',
                    'isDeleted': False
                }

            def merge_or_append_civitai(civitai_entry, preserve_existing_weight=False):
                civitai_id = get_version_id(civitai_entry)
                civitai_hash = (civitai_entry.get('hash') or '').lower()
                for index, existing in enumerate(loras):
                    existing_id = get_version_id(existing)
                    existing_hash = (existing.get('hash') or '').lower()
                    if not (
                        (civitai_id and existing_id == civitai_id)
                        or (civitai_hash and existing_hash == civitai_hash)
                    ):
                        continue

                    if preserve_existing_weight:
                        civitai_entry['weight'] = existing.get('weight', civitai_entry['weight'])
                    existing_base = existing.get('baseModel')
                    if not civitai_entry.get('baseModel'):
                        civitai_entry['baseModel'] = existing_base or ''
                    elif existing_base:
                        remaining = base_model_counts.get(existing_base, 0) - 1
                        if remaining > 0:
                            base_model_counts[existing_base] = remaining
                        else:
                            base_model_counts.pop(existing_base, None)
                    loras[index] = civitai_entry
                    return
                loras.append(civitai_entry)

            def merge_or_append_local(local_entry):
                local_id = get_version_id(local_entry)
                local_hash = (local_entry.get('hash') or '').lower()
                for existing in loras:
                    existing_id = get_version_id(existing)
                    existing_hash = (existing.get('hash') or '').lower()
                    if not (
                        (local_id and existing_id == local_id)
                        or (local_hash and existing_hash == local_hash)
                    ):
                        continue

                    existing['weight'] = local_entry['weight']
                    existing['hash'] = local_entry['hash']
                    existing['file_name'] = local_entry['file_name']
                    existing['existsLocally'] = True
                    existing['localPath'] = local_entry['localPath']
                    existing['size'] = local_entry['size']
                    existing['isDeleted'] = False
                    if not existing.get('modelId') and local_entry.get('modelId'):
                        existing['modelId'] = local_entry['modelId']
                    if not existing.get('baseModel') and local_entry.get('baseModel'):
                        existing['baseModel'] = local_entry['baseModel']
                        base_model_counts[local_entry['baseModel']] = base_model_counts.get(local_entry['baseModel'], 0) + 1
                    thumbnail_url = local_entry.get('thumbnailUrl')
                    if thumbnail_url and not thumbnail_url.endswith('/images/no-preview.png'):
                        existing['thumbnailUrl'] = thumbnail_url
                    return

                if local_entry.get('baseModel'):
                    base_model = local_entry['baseModel']
                    base_model_counts[base_model] = base_model_counts.get(base_model, 0) + 1
                loras.append(local_entry)

            resolved_prompt_basenames = set()
            queried_local_basenames = set()
            for lora_type, lora_name, lora_weight in prompt_loras.values():
                weight = round(float(lora_weight), 2)
                basename_key = (lora_type, normalize_lora_name(lora_name, True))
                matching_resources = [
                    lora
                    for lora in loras[:resource_lora_count]
                    if lora.get('file_name')
                    and normalize_lora_name(lora['file_name'], True) == basename_key[1]
                    and (
                        (lora_type == 'hypernet' and str(lora.get('type', '')).casefold() in ('hypernet', 'hypernetwork'))
                        or (lora_type == 'lora' and str(lora.get('type', '')).casefold() not in ('hypernet', 'hypernetwork'))
                    )
                ]
                if len(prompt_by_basename[basename_key]) == 1 and len(matching_resources) == 1:
                    matching_resources[0]['weight'] = weight
                    if basename_key not in hash_basenames:
                        resolved_prompt_basenames.add(basename_key)
                    continue

                if basename_key in hash_basenames:
                    continue

                if not recipe_scanner or lora_type != 'lora':
                    continue
                queried_local_basenames.add(basename_key)
                local_lora = await recipe_scanner.get_local_lora(lora_name, recipe_base_model)
                if not local_lora:
                    continue

                local_entry = self.populate_lora_from_local(
                    make_lora_entry(lora_type, lora_name, weight),
                    local_lora,
                )
                merge_or_append_local(local_entry)
                resolved_prompt_basenames.add(basename_key)

            for hash_key, lora_hash in metadata.get("hashes", {}).items():
                if not hash_key.startswith(("lora:", "hypernet:")):
                    continue
                lora_type, lora_name = hash_key.split(':', 1)
                basename_key = (lora_type, normalize_lora_name(lora_name, True))
                if basename_key in resolved_prompt_basenames:
                    continue

                prompt_entries = prompt_by_basename.get(basename_key, [])
                weight = prompt_entries[0][1] if len(prompt_entries) == 1 else 1.0
                lora_entry = make_lora_entry(lora_type, lora_name, weight, lora_hash)

                if lora_hash and recipe_scanner and lora_type == 'lora':
                    local_lora = await recipe_scanner.get_local_lora_by_hash(lora_hash)
                    if local_lora:
                        local_entry = self.populate_lora_from_local(lora_entry, local_lora)
                        merge_or_append_local(local_entry)
                        continue

                hash_resolved = False
                if lora_hash and metadata_provider:
                    try:
                        civitai_info = await metadata_provider.get_model_by_hash(lora_hash)
                        populated_entry = await self.populate_lora_from_civitai(
                            lora_entry,
                            civitai_info,
                            recipe_scanner,
                            base_model_counts,
                            lora_hash,
                        )
                        if populated_entry is None:
                            continue
                        lora_entry = populated_entry
                        hash_resolved = not lora_entry.get('isDeleted')
                    except Exception as e:
                        logger.error(f"Error fetching Civitai info for LoRA {lora_name}: {e}")

                if hash_resolved:
                    merge_or_append_civitai(lora_entry, preserve_existing_weight=not prompt_entries)
                    continue

                if recipe_scanner and lora_type == 'lora' and basename_key not in queried_local_basenames:
                    local_lora = await recipe_scanner.get_local_lora(lora_name, recipe_base_model)
                    if local_lora:
                        local_entry = self.populate_lora_from_local(lora_entry, local_lora)
                        merge_or_append_local(local_entry)
                        continue

                if lora_hash and not resource_lora_count:
                    loras.append(lora_entry)
                
            # Try to get base model from resources or make educated guess
            base_model = None
            if checkpoint and checkpoint.get("baseModel"):
                base_model = checkpoint.get("baseModel")
            elif base_model_counts:
                # Use the most common base model from the loras
                base_model = max(base_model_counts.items(), key=lambda x: x[1])[0]
            
            # Prepare final result structure
            # Make sure gen_params only contains recognized keys
            filtered_gen_params = {}
            for key in GEN_PARAM_KEYS:
                if key in metadata.get("gen_params", {}):
                    filtered_gen_params[key] = metadata["gen_params"][key]
            
            result = {
                'base_model': base_model,
                'loras': loras,
                'gen_params': filtered_gen_params,
                'from_automatic_metadata': True
            }

            if checkpoint:
                result['checkpoint'] = checkpoint
                result['model'] = checkpoint
            
            return result
            
        except Exception as e:
            logger.error(f"Error parsing Automatic1111 metadata: {e}", exc_info=True)
            return {"error": str(e), "loras": []}