#!/usr/bin/env python3
"""
Simplified flow registration script for self-hosted Prefect server.
Replaces Jenkins-based registration with direct registration to EA Prefect server.
"""

import prefect
import prefect.utilities.storage
import click
import importlib
import os
import sys
from pathlib import Path

# Storage and Run Config for self-hosted Prefect
from prefect.storage import Local, GitHub, Git
from prefect.run_configs import KubernetesRun
from prefect import Flow

# Add flows directory to Python path
sys.path.append(os.path.join(os.path.dirname(os.path.dirname(__file__)), "flows"))

def has_flow(file_path):
    """Check if a Python file contains a Prefect Flow."""
    try:
        dir_path = os.path.dirname(file_path)
        filename = os.path.basename(file_path)
        
        if not filename.endswith('.py'):
            return False
            
        # Add directory to path temporarily
        if dir_path not in sys.path:
            sys.path.insert(0, dir_path)
            
        module_name = filename.replace(".py", "")
        module = importlib.import_module(module_name)
        
        # Check if module contains any Flow objects
        for _, obj in module.__dict__.items():
            if isinstance(obj, prefect.Flow):
                return True
        return False
    except Exception as e:
        print(f"Error checking file {file_path}: {e}")
        return False


@click.command()
@click.option('--project', required=True, help='Prefect project name')
@click.option('--file', 'flow_file', required=True, help='Path to flow file')
@click.option('--server-url', default='https://ea-prefectv1-apollo-dev.int.enverus.com/graphql', 
              help='Prefect server GraphQL URL')
@click.option('--storage-type', default='git', type=click.Choice(['local', 'git']),
              help='Storage type: local or git')
@click.option('--git-repo', default='https://github.com/enverus-ea/ea.data.mfg.prefectV1.git', 
              help='Git repository URL')
@click.option('--git-ref', default='main', help='Git reference (branch/tag)')
@click.option('--labels', multiple=True, help='Flow labels')
@click.option('--image', default='392865356492.dkr.ecr.us-east-1.amazonaws.com/eashared-prefectv1-dev:pf-12176ef',
              help='Container image for flow execution (EA agent image)')
@click.option('--namespace', default='prefect', help='Kubernetes namespace')
def register_flow(project, flow_file, server_url, storage_type, git_repo, git_ref, labels, image, namespace):
    """Register a single flow with the self-hosted Prefect server."""
    
    print(f"Registering flow: {flow_file}")
    print(f"Project: {project}")
    print(f"Server: {server_url}")
    print(f"Storage: {storage_type}")
    print(f"Image: {image}")
    
    # Set Prefect server URL
    os.environ['PREFECT__CLOUD__API'] = server_url
    
    if not os.path.exists(flow_file):
        print(f"Flow file does not exist: {flow_file}")
        return
    
    # Add necessary paths to Python path before extracting flow
    paths_to_add = [
        os.path.dirname(flow_file),  # Directory containing the flow file
        os.path.dirname(os.path.dirname(__file__)),  # Root directory (where tasks folder is)
        os.path.join(os.path.dirname(os.path.dirname(__file__)), "flows"),
        ".",  # Current directory
    ]
    
    # Debug: Print the paths being added
    print("Adding paths to Python path:")
    for path in paths_to_add:
        if path and os.path.exists(path):
            print(f"  - {os.path.abspath(path)}")
            if path not in sys.path:
                sys.path.insert(0, path)
        else:
            print(f"  - {path} (does not exist)")
    
    # Debug: Check if tasks directory is accessible
    tasks_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "tasks")
    print(f"Tasks directory exists: {os.path.exists(tasks_path)}")
    if os.path.exists(tasks_path):
        print(f"Tasks directory contents: {os.listdir(tasks_path)}")


    # Configure storage
    if storage_type == 'git':
        # enverus-ea/ea.data.mfg.prefectV1
        # storage = GitHub(
        #     repo=git_repo,
        #     path=flow_file,
        #     ref=git_ref,
        #     access_token_secret="GITHUB_ACCESS_TOKEN"
        # )
        # full repo
        storage = Git(
            repo=git_repo,
            repo_host="github.com",
            flow_path=flow_file,
            branch_name=git_ref,
            git_token_secret_name="GITHUB_ACCESS_TOKEN"
        )
    else:
        # Use Local storage with absolute paths
        flow_dir = os.path.dirname(os.path.abspath(flow_file))
        flow_name = os.path.basename(flow_file)
    
        print(f"Local storage - Directory: {flow_dir}")
        print(f"Local storage - File: {flow_name}")
    
        storage = Local(
            directory=flow_dir if flow_dir else ".",
            path=flow_name
        )
    
    job_template = {
        "apiVersion": "batch/v1",
        "kind": "Job",
        "spec": {
            "template": {
                "spec": {
                    "nodeSelector": {
                        "kubernetes.io/arch": "amd64"
                    },
                    "restartPolicy": "Never",
                    "containers": [{
                        "name": "flow",
                        "env": [
                        {
                            "name": "PREFECT__CONTEXT__SECRETS__GITHUB_ACCESS_TOKEN",
                            "valueFrom": {
                                "secretKeyRef": {
                                    "name": "github-credentials",
                                    "key": "github-access-token",
                                    }
                                }
                            }
                        ]
                    }]
                }
            }
        }
    }
    
    # Configure Kubernetes run config with the EA agent image
    run_config = KubernetesRun(
        image=image,
        labels=list(labels) if labels else [],
        memory_request="512Mi",
        memory_limit="1Gi",
        cpu_request="250m",
        job_template=job_template
    )
    
    try:
        # Skip actual flow extraction during registration - just create a dummy flow
        # The real flow will be loaded from Git storage when it actually runs
        
        # Create a minimal flow object for registration
        flow_name = os.path.basename(flow_file).replace('.py', '')
        flow = Flow(name=flow_name)
        flow.storage = storage
        flow.run_config = run_config
        
        flow_id = flow.register(project_name=project)
        print(f"Successfully registered flow: {flow.name} (ID: {flow_id})")
        
    except Exception as e:
        print(f"Error registering flow {flow_file}: {e}")
        raise


@click.command()
@click.option('--project', required=True, help='Prefect project name')
@click.option('--pattern', required=True, help='Flow file pattern to search for (e.g., "permit.py" or "bdd/flows/permit.py")')
@click.option('--server-url', default='https://ea-prefectv1-apollo-dev.int.enverus.com/graphql',
              help='Prefect server GraphQL URL')
@click.option('--storage-type', default='git', type=click.Choice(['local', 'git']),
              help='Storage type: local or git')
@click.option('--git-repo', default='https://github.com/enverus-ea/ea.data.mfg.prefectV1.git',
              help='Git repository URL')
@click.option('--git-ref', default='main', help='Git reference (branch/tag)')
@click.option('--labels', multiple=True, help='Flow labels')
@click.option('--image', default='392865356492.dkr.ecr.us-east-1.amazonaws.com/eashared-prefectv1-dev:pf-12176ef',
              help='Container image for flow execution (EA agent image)')
@click.option('--namespace', default='prefect', help='Kubernetes namespace')
def find_and_register(pattern, project, server_url, storage_type, git_repo, git_ref, labels, image, namespace):
    """Find and register flows matching a pattern."""
    
    print(f"Searching for flows matching pattern: {pattern}")
    
    # Search for files matching the pattern
    matching_files = []
    
    # If pattern is a full path, check if it exists
    if os.path.exists(pattern):
        matching_files = [pattern]
    else:
        # Search recursively for files matching the pattern
        for root, dirs, files in os.walk('.'):
            for file in files:
                if file.endswith('.py'):
                    full_path = os.path.join(root, file)
                    # Check if pattern matches filename or full path
                    if pattern in full_path or pattern == file:
                        matching_files.append(full_path)
    
    if not matching_files:
        print(f"No flows found matching pattern: {pattern}")
        return
    
    print(f"Found {len(matching_files)} matching files:")
    for file in matching_files:
        print(f"  - {file}")
    
    # Register each matching flow
    for flow_file in matching_files:
        try:
            register_flow.callback(
                project=project,
                flow_file=flow_file,
                server_url=server_url,
                storage_type=storage_type,
                git_repo=git_repo,
                git_ref=git_ref,
                labels=labels,
                image=image,
                namespace=namespace
            )
        except Exception as e:
            print(f"Failed to register {flow_file}: {e}")
            continue


@click.group()
def cli():
    """Flow registration utilities for self-hosted Prefect."""
    pass

cli.add_command(find_and_register, name='find')
cli.add_command(register_flow, name='single')


if __name__ == "__main__":
    cli()