#!/bin/env python3
import requests
import os
import glob

def create_knowledge_base_from_repo(token, base_url, repo_path, kb_name, kb_purpose="Documentation"):
    """
    Create a knowledge base from a git repository
    
    Args:
        token (str): Your OpenWebUI API token
        base_url (str): Base URL of your OpenWebUI instance (e.g. http://localhost:3000)
        repo_path (str): Path to the local git repository
        kb_name (str): Name for the knowledge base
        kb_purpose (str): Purpose description for the knowledge base
    """
    headers = {
        'Authorization': f'Bearer {token}',
        'Accept': 'application/json'
    }

    # 1. Create knowledge base
    kb_url = f"{base_url}/api/v1/knowledge/create"
    kb_data = {
        "name": kb_name,
        "description": kb_purpose
    }
    kb_response = requests.post(kb_url, headers=headers, json=kb_data)
    if not kb_response.ok:
        raise Exception(f"Failed to create knowledge base: {kb_response.text}")
    
    kb_id = kb_response.json().get('id')
    print(f"Created knowledge base with ID: {kb_id}")

    # 2. Upload files and add them to knowledge base
    file_types = ['*.txt', '*.md', '*.py', '*.js', '*.html', '*.css', '*.json', '*.yaml', '*.yml', '*.sh', '*.bat', 'Dockerfile', 'LICENSE', '*.tsx', '*.j2', '*.tpl', '*.ts' ]
    files_processed = 0

    for file_pattern in file_types:
        for file_path in glob.glob(os.path.join(repo_path, '**', file_pattern), recursive=True):
            try:
                # Skip .git directory
                if '.git' in file_path:
                    continue
                    
                # Upload file
                with open(file_path, 'rb') as file:
                    files = {'file': file}
                    upload_response = requests.post(
                        f"{base_url}/api/v1/files/",
                        headers=headers,
                        files=files
                    )
                    
                if not upload_response.ok:
                    print(f"Failed to upload {file_path}: {upload_response.text}")
                    continue

                file_id = upload_response.json().get('id')

                # Add file to knowledge base
                add_file_response = requests.post(
                    f"{base_url}/api/v1/knowledge/{kb_id}/file/add",
                    headers={'Authorization': f'Bearer {token}', 'Content-Type': 'application/json'},
                    json={'file_id': file_id}
                )

                if add_file_response.ok:
                    files_processed += 1
                    print(f"Successfully processed: {file_path}")
                else:
                    print(f"Failed to add {file_path} to knowledge base: {add_file_response.text}")

            except Exception as e:
                print(f"Error processing {file_path}: {str(e)}")

    print(f"\nProcessing complete! Added {files_processed} files to knowledge base {kb_name}")

if __name__ == "__main__":
    import argparse

    # Set up argument parser
    parser = argparse.ArgumentParser(description='Create a knowledge base from a git repository')
    parser.add_argument('--token', '-t', required=True, help='Your OpenWebUI API token')
    parser.add_argument('--base-url', '-u', required=True, help='Base URL of your OpenWebUI instance')
    parser.add_argument('--repo-path', '-r', required=True, help='Path to the local git repository')
    parser.add_argument('--kb-name', '-n', required=True, help='Name for the knowledge base')
    parser.add_argument('--kb-purpose', '-p', required=False, help='Purpose/description for the knowledge base')
    
    # Parse arguments
    args = parser.parse_args()
    
    # Build kwargs dictionary
    kwargs = {
        'token': args.token,
        'base_url': args.base_url,
        'repo_path': args.repo_path,
        'kb_name': args.kb_name
    }
    
    # Only add kb_purpose if it was provided
    if args.kb_purpose:
        kwargs['kb_purpose'] = args.kb_purpose

    # Use parsed arguments
    create_knowledge_base_from_repo(**kwargs)