#!/usr/bin/env python3
import os
import sys
import json
from pathlib import Path
from kudra_cloud_client import KudraCloudClient

KUDRA_TOKEN = os.getenv("KUDRA_TOKEN")
KUDRA_PROJECT_RUN_ID = os.getenv("KUDRA_PROJECT_RUN_ID")

if not KUDRA_TOKEN or not KUDRA_PROJECT_RUN_ID:
    print(json.dumps({"error": "Missing KUDRA_TOKEN or KUDRA_PROJECT_RUN_ID environment variables"}))
    sys.exit(1)

def extract_documents(path: str):
    """Extract structured data from document(s)"""
    
    # Determine if path is file or directory
    if os.path.isfile(path):
        file_dir = str(Path(path).parent)
    elif os.path.isdir(path):
        file_dir = path
    else:
        return {"error": f"Path not found: {path}"}
    
    # Initialize Kudra client
    kudra_client = KudraCloudClient(token=KUDRA_TOKEN)
    
    try:
        # Process documents
        results = kudra_client.analyze_documents(
            files_dir=file_dir,
            project_run_id=KUDRA_PROJECT_RUN_ID
        )
        
        if not results:
            return {"error": "No results returned from Kudra"}
        
        # Format output
        output = []
        for result in results:
            doc_data = {
                "filename": result.get('file', 'Unknown'),
                "text_length": len(result.get('text', '')),
                "tables_found": len(result.get('extracted_tables', [])),
                "has_metadata": bool(result.get('open_ai_result'))
            }
            
            # Include first table as sample
            if result.get('extracted_tables'):
                table = result['extracted_tables'][0]
                doc_data['sample_table'] = {
                    "rows": len(table.get('data', [])),
                    "preview": str(table.get('data', [])[:2])
                }
            
            output.append(doc_data)
        
        return {"success": True, "documents": output, "full_results": results}
        
    except Exception as e:
        return {"error": str(e)}

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(json.dumps({"error": "Usage: extract.py <file_or_directory_path>"}))
        sys.exit(1)
    
    path = sys.argv[1]
    result = extract_documents(path)
    print(json.dumps(result, indent=2))