#!/usr/bin/env python3
"""fd-thallium-app"""

__version__ = '0.4.3'

import argparse
import json
import logging
import os
import re
import fnmatch
import sys
from pathlib import Path
from typing import Optional, List, Dict, Any, Union

import yaml
import hvac
from dotenv import load_dotenv

SYSLOG_NAME = "fd-thallium-app"
LOG = logging.getLogger(SYSLOG_NAME)

def parse_args() -> argparse.Namespace:
    """Parse and validate command line arguments"""
    parser = argparse.ArgumentParser(description="fd-thallium-app")
    
    parser.add_argument("-l", "--loglevel",
                       default='info',
                       choices=['critical', 'error', 'warning', 'info', 'debug'],
                       help="Emit traces with LOGLEVEL details")
    
    parser.add_argument("-c", "--conffile",
                       default='~/.config/fd-thallium',
                       help="Use configuration file <conffile>")
    
    parser.add_argument("--uri", "--thallium-uri",
                       dest='uri',
                       help="Thallium URI address")
    
    parser.add_argument("--role-id", "--thallium-role-id",
                       dest='role_id',
                       help="Thallium role ID")
    
    parser.add_argument("--secret-id", "--thallium-secret-id",
                       dest='secret_id',
                       help="Thallium secret ID")
    
    parser.add_argument("--glob",
                       action='store_true',
                       default=False,
                       help="Allow glob in key name for modes list and read")
    
    parser.add_argument("--indent",
                       type=int,
                       default=None,
                       help="Indent level for output")
    
    parser.add_argument("--key-name", "--thallium-key-name",
                       dest='key_name',
                       help="Thallium key name")
    
    parser.add_argument("--token", "--thallium-token",
                       dest='token',
                       help="Thallium token")
    
    parser.add_argument("--format",
                       default='raw',
                       choices=['raw', 'json', 'yaml'],
                       help="Output format")
    
    parser.add_argument("-o", "--output", "--data",
                       default='-',
                       help="Output data in file or to stdout")
    
    parser.add_argument("--metatype",
                       default='data',
                       choices=['data', 'metadata'],
                       help="Output metatype")
    
    parser.add_argument("-m", "--mode",
                       default='read',
                       choices=['read', 'rmerge', 'write', 'wmerge', 'merge', 'list'],
                       help="Operation mode")
    
    parser.add_argument("-k", "--keys",
                       action='append',
                       default=[],
                       help="Extract keys in file or to stdout")

    args = parser.parse_args()
    args.loglevel = getattr(logging, args.loglevel.upper(), logging.INFO)
    return args

class FdThalliumApp:
    def __init__(self, options: argparse.Namespace):
        self.options = options
        self.prevdata: Dict = {}
        self.merging: bool = False
        self.client: Optional[hvac.Client] = None
        
        self._load_config()
        self._init_client()

    def _load_config(self) -> None:
        """Load configuration from file and environment variables"""
        if self.options.conffile:
            config_path = Path(self.options.conffile).expanduser()
            if config_path.exists():
                load_dotenv(dotenv_path=config_path)

        self.options.uri = self.options.uri or os.getenv('THALLIUM_URI')
        if not self.options.uri:
            raise ValueError("missing variable THALLIUM_URI")

        self.options.token = self.options.token or os.getenv('THALLIUM_TOKEN') or os.getenv('VAULT_TOKEN')
        if not self.options.token or self.options.token == '-':
            self.options.role_id = self.options.role_id or os.getenv('THALLIUM_ROLE_ID')
            self.options.secret_id = self.options.secret_id or os.getenv('THALLIUM_SECRET_ID')
            
            if not self.options.role_id:
                raise ValueError("missing variable THALLIUM_ROLE_ID")
            if not self.options.secret_id:
                raise ValueError("missing variable THALLIUM_SECRET_ID")

        self.options.key_name = self.options.key_name or os.getenv('THALLIUM_KEY_NAME')
        if not self.options.key_name:
            raise ValueError("missing variable THALLIUM_KEY_NAME")

    def _init_client(self) -> None:
        """Initialize hvac client with appropriate authentication"""
        try:
            if not self.options.uri:
                raise ValueError("Vault URI not specified")

            self.client = hvac.Client(url=self.options.uri)
            
            if self.options.token:
                self.client.token = self.options.token
            elif self.options.role_id and self.options.secret_id:
                auth_response = self.client.auth.approle.login(
                    role_id=self.options.role_id,
                    secret_id=self.options.secret_id
                )
                self.client.token = auth_response['auth']['client_token']
            else:
                raise ValueError("Either token or role_id/secret_id must be provided")
                
            if not self.client.is_authenticated():
                raise hvac.exceptions.VaultError("Failed to authenticate with Vault")
                
        except (hvac.exceptions.VaultError, ValueError) as e:
            LOG.error(f"Failed to initialize Vault client: {e}")
            raise

    def _output(self, data: Union[str, Dict, List]) -> None:
        """Format and output data according to specified format"""
        if not isinstance(data, str):
            if self.options.format == 'yaml':
                data = yaml.safe_dump(data, default_flow_style=False)
            else:
                data = json.dumps(data, indent=self.options.indent)
        elif self.options.format in ('yaml', 'json'):
            try:
                parsed_data = json.loads(data) if self.options.format == 'json' else yaml.safe_load(data)
                data = (json.dumps(parsed_data, indent=self.options.indent) 
                       if self.options.format == 'json' 
                       else yaml.safe_dump(parsed_data, default_flow_style=False))
            except (json.JSONDecodeError, yaml.YAMLError):
                pass  # Keep original format if parsing fails

        if self.options.output == '-':
            sys.stdout.write(data)
            sys.stdout.write('\n')
        else:
            Path(self.options.output).write_text(data, encoding='utf-8')

    def _list_filter(self, xpath: str, regex) -> Optional[List[str]]:
        """Filter and list secrets at given path"""
        try:
            response = self.client.secrets.kv.v2.list_secrets(
                path=xpath.lstrip('/'),
                mount_point='secret'
            )
        except hvac.exceptions.InvalidPath:
            return None

        if not response or 'data' not in response or 'keys' not in response['data']:
            return None

        return [f"{xpath}/{key.rstrip('/')}" 
                for key in response['data']['keys'] 
                if regex(key.rstrip('/'))]

    def _list_walk(self, xpath: str, regexs: List[Dict]) -> List[str]:
        """Recursively walk through paths matching regex patterns"""
        def _walk(wpath: str, regex: Dict, ret: List[str], pos: int = 0) -> None:
            xdirs = self._list_filter(wpath, regex['regex'])
            if not xdirs and pos < len(regexs):
                return

            if not xdirs or pos >= len(regexs):
                ret.append(wpath)
                return

            for xdir in xdirs:
                _walk(xdir, regexs[pos], ret, pos + 1)

        result: List[str] = []
        _walk(xpath, regexs[0], result)
        return result

    def do_list(self) -> None:
        """List secrets according to specified pattern"""
        if not self.options.glob:
            self._do_list_noglob()
            return

        regexs = []
        maxpath = []

        parts = [p for p in self.options.key_name.split('/') if p]
        
        # Find the first part containing a wildcard
        first_wildcard_idx = next((i for i, part in enumerate(parts) 
                                if '*' in part or '?' in part), len(parts))
        
        # Add all parts before the first wildcard to maxpath
        maxpath.extend(parts[:first_wildcard_idx])
        
        # Convert remaining parts to regex patterns
        for part in parts[first_wildcard_idx:]:
            trans = fnmatch.translate(part)
            # Remove the \Z from the end of the pattern
            trans = trans.replace('\Z', '')
            regexs.append({
                'translate': trans,
                'regex': re.compile(trans).match,
                'raw': part
            })

        if not regexs:
            self._do_list_noglob()
            return

        base_path = '/'.join(maxpath) if maxpath else ''
        result = self._list_walk(base_path, regexs)
        self._output(result)

    def _do_list_noglob(self) -> None:
        """List secrets without glob pattern matching"""
        try:
            response = self.client.secrets.kv.v2.list_secrets(
                path=self.options.key_name.lstrip('/'),
                mount_point='secret'
            )
        except hvac.exceptions.InvalidPath:
            self._output([])
            return

        result = []
        xpath = self.options.key_name.rstrip('/')

        if response and 'data' in response:
            if 'keys' in response['data']:
                result.extend(f"{xpath}/{key}" for key in response['data']['keys'])
            if 'data' in response['data']:
                result.extend(f"{xpath}/{key}" for key in response['data']['data'])

        self._output(result)

    def do_read(self, key_name: Optional[str] = None) -> Dict[str, Any]:
        """Read secret from specified path"""
        key_name = key_name or self.options.key_name
        try:
            response = self.client.secrets.kv.v2.read_secret_version(
                path=key_name.lstrip('/'),
                mount_point='secret',
                raise_on_deleted_version=True
            )
            
            if response and 'data' in response:
                secret_data = response['data']['data']
                
                # If specific keys are requested and not called from another method
                if self.options.keys and not self.merging:
                    for key in self.options.keys:
                        # Handle key=file syntax
                        if '=' in key:
                            key_name, file_path = key.split('=', 1)
                            if key_name in secret_data:
                                value = str(secret_data[key_name])
                                if file_path == '-':
                                    sys.stdout.write(value)
                                    sys.stdout.write('\n')
                                else:
                                    Path(file_path).write_text(value + '\n', encoding='utf-8')
                        # Handle just key (output to stdout)
                        elif key in secret_data:
                            value = str(secret_data[key])
                            sys.stdout.write(value)
                            sys.stdout.write('\n')
                elif not self.merging:
                    # Output entire secret if no specific keys requested
                    self._output(secret_data)
                
                return secret_data
            else:
                if not self.merging:
                    self._output({})
                return {}
                
        except hvac.exceptions.VaultError as e:
            LOG.error(f"Failed to read secret: {e}")
            raise

    def do_write(self) -> None:
        """Write secret to specified path"""
        try:
            # Read input data
            if self.options.output == '-':
                data = json.load(sys.stdin)
            else:
                data = json.loads(Path(self.options.output).read_text(encoding='utf-8'))

            # Write to Vault
            self.client.secrets.kv.v2.create_or_update_secret(
                path=self.options.key_name.lstrip('/'),
                secret=data,
                mount_point='secret'
            )
            
            # Output the written data if needed
            self._output(data)
        except json.JSONDecodeError as e:
            LOG.error(f"Invalid JSON input: {e}")
            raise
        except hvac.exceptions.VaultError as e:
            LOG.error(f"Failed to write secret: {e}")
            raise

    def do_merge(self) -> None:
        """Merge new data with existing secret"""
        try:
            # Read existing data
            try:
                existing_data = self.do_read(self.options.key_name)
                if not existing_data:
                    existing_data = {}
            except hvac.exceptions.InvalidPath:
                existing_data = {}

            # Read new data
            if self.options.output == '-':
                new_data = json.load(sys.stdin)
            else:
                new_data = json.loads(Path(self.options.output).read_text(encoding='utf-8'))

            # Merge data (new_data takes precedence)
            merged_data = {**existing_data, **new_data}

            # Write merged data back
            self.client.secrets.kv.v2.create_or_update_secret(
                path=self.options.key_name.lstrip('/'),
                secret=merged_data,
                mount_point='secret'
            )

            # Output the merged data
            self._output(merged_data)
        except json.JSONDecodeError as e:
            LOG.error(f"Invalid JSON input: {e}")
            raise
        except hvac.exceptions.VaultError as e:
            LOG.error(f"Failed to merge secret: {e}")
            raise
    
    def do_rmerge(self) -> None:
        """Read and merge multiple secrets"""
        try:
            merged_data = {}
            self.merging = True

            # If glob is enabled, get all matching paths
            if self.options.glob:
                regexs = []
                maxpath = []

                parts = [p for p in self.options.key_name.split('/') if p]
                first_wildcard_idx = next((i for i, part in enumerate(parts)
                                           if '*' in part or '?' in part), len(parts))
                maxpath.extend(parts[:first_wildcard_idx])

                for part in parts[first_wildcard_idx:]:
                    trans = fnmatch.translate(part).replace('\\Z', '')
                    regexs.append({
                        'translate': trans,
                        'regex': re.compile(trans).match,
                        'raw': part
                    })

                base_path = '/'.join(maxpath) if maxpath else ''
                paths = self._list_walk(base_path, regexs)
            else:
                paths = [self.options.key_name]

            # Read and merge all secrets
            for path in paths:
                try:
                    secret_data = self.do_read(path)
                    if self.options.keys and secret_data:
                        for key in self.options.keys:
                            if key in secret_data:
                                value = secret_data[key]
                                if isinstance(value, str) and value.strip().startswith('{'):
                                    try:
                                        parsed_data = json.loads(value)
                                        merged_data.update(parsed_data)
                                    except json.JSONDecodeError:
                                        LOG.warning(f"Invalid JSON content in {key}")
                except hvac.exceptions.InvalidPath:
                    LOG.warning(f"Path not found: {path}")
                    continue

            self._output(merged_data)
        except hvac.exceptions.VaultError as e:
            LOG.error(f"Failed to read and merge secrets: {e}")
            raise

    def do_wmerge(self) -> None:
        """Write and merge data while preserving existing structure"""
        try:
            self.merging = True
            # Read new data
            if self.options.output == '-':
                new_data = json.load(sys.stdin)
            else:
                new_data = json.loads(Path(self.options.output).read_text(encoding='utf-8'))

            # Read existing data
            try:
                existing_data = self.do_read(self.options.key_name)
                if not existing_data:
                    existing_data = {}
            except hvac.exceptions.InvalidPath:
                existing_data = {}

            # Deep merge function
            def deep_merge(d1, d2):
                merged = d1.copy()
                for k, v in d2.items():
                    if k in merged and isinstance(merged[k], dict) and isinstance(v, dict):
                        merged[k] = deep_merge(merged[k], v)
                    else:
                        merged[k] = v
                return merged

            # If a specific key is provided, handle it specially
            if self.options.keys:
                key = self.options.keys[0]
                if key in existing_data:
                    # If the existing value is a JSON string, parse it
                    existing_value = existing_data[key]
                    if isinstance(existing_value, str) and existing_value.strip().startswith('{'):
                        try:
                            existing_json = json.loads(existing_value)
                            merged_json = deep_merge(existing_json, new_data)
                            existing_data[key] = json.dumps(merged_json, indent=2)
                        except json.JSONDecodeError:
                            existing_data[key] = json.dumps(new_data, indent=2)
                    else:
                        existing_data[key] = json.dumps(new_data, indent=2)
                else:
                    existing_data[key] = json.dumps(new_data, indent=2)
                merged_data = existing_data
            else:
                # If no specific key, merge at top level as before
                merged_data = deep_merge(existing_data, new_data)

            # Write merged data back
            self.client.secrets.kv.v2.create_or_update_secret(
                path=self.options.key_name.lstrip('/'),
                secret=merged_data,
                mount_point='secret'
            )

            # Output the merged data
            self._output(merged_data)
        except json.JSONDecodeError as e:
            LOG.error(f"Invalid JSON input: {e}")
            raise
        except hvac.exceptions.VaultError as e:
            LOG.error(f"Failed to write merged secret: {e}")
            raise

def main() -> None:
    """Main entry point"""
    try:
        options = parse_args()
        logging.basicConfig(level=options.loglevel)
        
        app = FdThalliumApp(options)
        
        mode_handlers = {
            'list': app.do_list,
            'read': app.do_read,
            'write': app.do_write,
            'merge': app.do_merge,
            'rmerge': app.do_rmerge,
            'wmerge': app.do_wmerge
        }
        
        if options.mode in mode_handlers:
            mode_handlers[options.mode]()
        else:
            LOG.error(f"Unsupported mode: {options.mode}")
            sys.exit(1)
            
    except Exception as e:
        LOG.error(f"Error: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()