#!/usr/bin/env python3
import sys
import argparse
import paramiko
import os

def fetch_logs(args):
    ssh = paramiko.SSHClient()
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    
    try:
        # Connect to the server
        print(f"Connecting to {args.host}...")
        ssh.connect(
            hostname=args.host,
            username=args.username,
            password=args.password,
            timeout=10
        )
        
        command = ""
        if args.mode == "docker":
            since_flag = f'--since "{args.since}"' if args.since else ""
            until_flag = f'--until "{args.until}"' if args.until else ""
            tail_flag = f"--tail={args.lines}" if not args.since else ""
            # docker-compose logs command
            command = f"docker-compose -f {args.yml_path} logs --no-log-prefix {since_flag} {until_flag} {tail_flag} {args.service}"
        else:
            # File mode
            if args.since or args.until:
                # Time-based filtering for file mode using awk
                since_val = args.since if args.since else ""
                until_val = args.until if args.until else ""
                
                # Convert relative time (e.g., 10m) to absolute timestamp on server
                if since_val and any(since_val.endswith(suffix) for suffix in ['s', 'm', 'h', 'd']) and since_val[:-1].isdigit():
                    unit_map = {'s': 'seconds', 'm': 'minutes', 'h': 'hours', 'd': 'days'}
                    unit = unit_map[since_val[-1]]
                    num = since_val[:-1]
                    since_expr = f"\"$(date -d '{num} {unit} ago' '+%Y-%m-%d %H:%M:%S')\""
                elif since_val:
                    since_expr = f"'{since_val}'"
                else:
                    since_expr = "''"
                
                until_expr = f"'{until_val}'" if until_val else "''"
                
                # awk script that extracts timestamp and compares
                # It maintains 'in_range' state for multi-line logs (like stack traces)
                awk_script = (
                    f"awk -v since={since_expr} -v until={until_expr} '"
                    "BEGIN { in_range = (since == \"\"); } "
                    "{ "
                    "  if (match($0, /^[0-9]{4}-[0-9]{2}-[0-9]{2} [0-9]{2}:[0-9]{2}:[0-9]{2}/)) { "
                    "    current_time = substr($0, RSTART, 19); "
                    "    in_range = (since == \"\" || current_time >= since) && (until == \"\" || current_time <= until); "
                    "  } "
                    "  if (in_range) print $0; "
                    "}' " + args.log_path
                )
                command = awk_script
            else:
                # Default tail + grep for errors
                command = f"tail -n {args.lines} {args.log_path} | grep -E 'Exception|Error|Caused by' -C 50"
        
        print(f"Executing command: {command}")
        stdin, stdout, stderr = ssh.exec_command(command)
        
        output = stdout.read().decode('utf-8', errors='replace')
        error = stderr.read().decode('utf-8', errors='replace')
        
        if error and not output:
            print(f"Error from server:\n{error}")
            return
            
        if not output:
            print("No matching logs found.")
            return

        # Print a header for the output
        print("-" * 40)
        print(f"LOGS FROM {args.host} ({args.service if args.mode == 'docker' else args.log_path})")
        print("-" * 40)
        print(output)
        print("-" * 40)
        
    except Exception as e:
        print(f"Failed to fetch logs: {str(e)}")
    finally:
        ssh.close()

def main():
    parser = argparse.ArgumentParser(description="Fetch logs from a remote server via SSH.")
    parser.add_argument("--host", required=True, help="Server IP or hostname")
    parser.add_argument("--username", required=True, help="SSH username")
    parser.add_argument("--password", help="SSH password")
    parser.add_argument("--mode", choices=["docker", "file"], required=True, help="Log extraction mode")
    parser.add_argument("--service", help="Docker service name (for docker mode)")
    parser.add_argument("--yml-path", help="Path to docker-compose.yml (for docker mode)")
    parser.add_argument("--log-path", help="Path to log file (for file mode)")
    parser.add_argument("--lines", type=int, default=1000, help="Number of lines to fetch/analyze (default 1000)")
    parser.add_argument("--since", help="Filter logs since this time (e.g., '10m', '2026-01-16 10:10:00')")
    parser.add_argument("--until", help="Filter logs until this time (e.g., '2026-01-16 10:20:00')")
    
    args = parser.parse_args()
    
    if args.mode == "docker" and not (args.service and args.yml_path):
        parser.error("docker mode requires --service and --yml-path")
    if args.mode == "file" and not args.log_path:
        parser.error("file mode requires --log-path")
        
    fetch_logs(args)

if __name__ == "__main__":
    main()