#!/usr/bin/env python3
"""
Deploy the demo data generator function and scheduled workflow.

This script deploys a Cognite Function that generates test time series data
for data quality validation testing. It only deploys if demo_data.enabled
is set to true in settings.yaml.

Usage:
    python scripts/deploy_demo_data.py [--dry-run] [--force]
"""

import argparse
import hashlib
import sys
import time
import zipfile
from pathlib import Path

import yaml
from cognite.client.data_classes.iam import ClientCredentials
from cognite.client.data_classes.workflows import (
    FunctionTaskParameters,
    WorkflowDefinitionUpsert,
    WorkflowScheduledTriggerRule,
    WorkflowTask,
    WorkflowTriggerUpsert,
    WorkflowUpsert,
    WorkflowVersionUpsert,
)
from shared import create_client, get_repo_root, load_settings


def compute_function_hash(repo_root: Path, settings: dict) -> str:
    """Compute hash of the demo data function code."""
    demo_settings = settings.get("demo_data", {})
    func_settings = demo_settings.get("function", {})
    source_dir = repo_root / func_settings.get("source_dir", "generate_demo_data")

    hasher = hashlib.sha256()

    # Hash included files
    for filename in func_settings.get("include_files", []):
        file_path = source_dir / filename
        if file_path.exists():
            hasher.update(file_path.read_bytes())

    # Hash settings
    hasher.update(yaml.dump(demo_settings, sort_keys=True).encode())

    return hasher.hexdigest()[:16]


def get_deployed_function_hash(client, function_external_id: str) -> str | None:
    """Get the hash of the currently deployed function from its metadata."""
    try:
        func = client.functions.retrieve(external_id=function_external_id)
        if func and func.metadata:
            return func.metadata.get("code_hash")
        return None
    except Exception:
        return None


def create_function_zip(repo_root: Path, settings: dict, output_path: Path) -> None:
    """Create a zip file containing the function code."""
    demo_settings = settings.get("demo_data", {})
    func_settings = demo_settings.get("function", {})
    source_dir = repo_root / func_settings.get("source_dir", "generate_demo_data")

    with zipfile.ZipFile(output_path, "w", zipfile.ZIP_DEFLATED) as zipf:
        for filename in func_settings.get("include_files", []):
            file_path = source_dir / filename
            if file_path.exists():
                zipf.write(file_path, filename)
                print(f"  Added {filename}")


def deploy_function(client, repo_root: Path, settings: dict, force: bool = False, dry_run: bool = False) -> bool:
    """Deploy the demo data generator function."""
    demo_settings = settings.get("demo_data", {})
    func_settings = demo_settings.get("function", {})
    function_external_id = func_settings.get("external_id", "generate-demo-data")

    current_hash = compute_function_hash(repo_root, settings)
    print(f"Current code hash: {current_hash}")

    deployed_hash = get_deployed_function_hash(client, function_external_id)
    print(f"Deployed code hash: {deployed_hash or 'N/A (not deployed)'}")

    needs_deploy = force or (current_hash != deployed_hash)

    if not needs_deploy:
        print("Function is up to date, skipping deployment")
        return False

    if dry_run:
        print(f"[DRY RUN] Would deploy function with hash: {current_hash}")
        return True

    # Create zip file
    zip_path = repo_root / "demo_function_deploy.zip"
    print("\nCreating zip file...")
    create_function_zip(repo_root, settings, zip_path)
    print(f"Zip created: {zip_path.stat().st_size:,} bytes")

    try:
        # Get dataset ID if configured
        dataset_external_id = settings.get("shacl_rules", {}).get("dataset_external_id")
        dataset_id = None

        if dataset_external_id:
            try:
                dataset = client.data_sets.retrieve(external_id=dataset_external_id)
                if dataset:
                    dataset_id = dataset.id
                    print(f"  Using dataset: {dataset_external_id} (id: {dataset_id})")
            except Exception as e:
                print(f"  WARNING: Dataset {dataset_external_id} not found: {e}")
                print("  Uploading without dataset association")

        # Upload to CDF Files
        print("\nUploading to CDF Files...")
        uploaded_file = client.files.upload(
            path=str(zip_path),
            name=f"{function_external_id}.zip",
            external_id=f"{function_external_id}-code",
            data_set_id=dataset_id,
            overwrite=True,
        )
        print(f"Uploaded with file_id: {uploaded_file.id}")

        # Delete existing function if present
        print("\nChecking for existing function...")
        try:
            client.functions.delete(external_id=function_external_id)
            print("Deleted existing function")
        except Exception:
            print("No existing function to delete")

        # Deploy function
        print("\nDeploying function...")
        function = client.functions.create(
            name=func_settings.get("name", "generate-demo-data"),
            external_id=function_external_id,
            file_id=uploaded_file.id,
            description=func_settings.get("description", "Generate demo data for testing"),
            runtime=func_settings.get("runtime", "py311"),
            metadata={"code_hash": current_hash},
        )
        print(f"Function created with ID: {function.id}")

        # Wait for deployment
        print("\nWaiting for function to be ready...")
        wait_time = 0
        while True:
            func = client.functions.retrieve(external_id=function_external_id)
            status = func.status
            print(f"  Status: {status} ({wait_time}s)")

            if status == "Ready":
                print(f"\nFunction ready after {wait_time}s")
                break
            elif status == "Failed":
                raise Exception(f"Function deployment failed: {func.error}")

            time.sleep(10)
            wait_time += 10

            if wait_time > 600:
                raise Exception("Function deployment timeout (10 minutes)")

        return True

    finally:
        if zip_path.exists():
            zip_path.unlink()


def deploy_workflow(client, settings: dict, force: bool = False, dry_run: bool = False) -> tuple[str, bool]:
    """Deploy the scheduled workflow for demo data generation."""
    demo_settings = settings.get("demo_data", {})
    workflow_settings = demo_settings.get("workflow", {})
    workflow_external_id = workflow_settings.get("external_id", "dq-demo-data-generator")
    version = workflow_settings.get("version", "1")

    if dry_run:
        print(f"[DRY RUN] Would create workflow: {workflow_external_id}")
        return workflow_external_id, True

    # Get dataset ID for workflow (use SHACL rules dataset)
    dataset_id = None
    shacl_rules = settings.get("shacl_rules", {})
    if shacl_rules.get("dataset_external_id"):
        dataset = client.data_sets.retrieve(external_id=shacl_rules["dataset_external_id"])
        if dataset:
            dataset_id = dataset.id

    # Create workflow
    workflow = WorkflowUpsert(
        external_id=workflow_external_id,
        description="Demo data generator workflow for data quality testing",
        data_set_id=dataset_id,
    )
    client.workflows.upsert(workflow)

    # Build task parameters
    task_settings = workflow_settings.get("task", {})
    generation_settings = demo_settings.get("generation", {})

    task_data = {
        "space": generation_settings.get("space", "ts_dq_test"),
        "historical_days": generation_settings.get("historical_days", 365),
        "points_per_run": generation_settings.get("points_per_run", 1),
    }

    func_settings = demo_settings.get("function", {})
    function_external_id = func_settings.get("external_id", "generate-demo-data")

    task = WorkflowTask(
        external_id="generate-demo-data-task",
        name="Generate Demo Data",
        description="Generate test time series data for data quality validation",
        parameters=FunctionTaskParameters(external_id=function_external_id, data=task_data),
        retries=task_settings.get("retries", 1),
        timeout=task_settings.get("timeout", 300),
        on_failure=task_settings.get("on_failure", "abortWorkflow"),
    )

    workflow_definition = WorkflowDefinitionUpsert(tasks=[task], description="Demo data generator workflow")

    workflow_version = WorkflowVersionUpsert(
        workflow_external_id=workflow_external_id, version=version, workflow_definition=workflow_definition
    )

    client.workflows.versions.upsert(workflow_version)
    print(f"Created workflow: {workflow_external_id} v{version}")

    return workflow_external_id, True


def deploy_trigger(client, settings: dict, workflow_external_id: str, dry_run: bool = False) -> tuple[str, bool]:
    """Deploy scheduled trigger for demo data generation."""
    demo_settings = settings.get("demo_data", {})
    workflow_settings = demo_settings.get("workflow", {})
    schedule_settings = demo_settings.get("schedule", {})

    trigger_external_id = "dq-demo-data-trigger"
    cron_expression = schedule_settings.get("cron", "* * * * *")

    if dry_run:
        print(f"[DRY RUN] Would create trigger: {trigger_external_id}")
        return trigger_external_id, True

    # Get credentials for trigger
    creds = client.config.credentials
    trigger_credentials = ClientCredentials(
        client_id=creds.client_id,
        client_secret=creds.client_secret,
    )

    # Create scheduled trigger
    trigger_rule = WorkflowScheduledTriggerRule(cron_expression=cron_expression)

    trigger = WorkflowTriggerUpsert(
        external_id=trigger_external_id,
        trigger_rule=trigger_rule,
        workflow_external_id=workflow_external_id,
        workflow_version=workflow_settings.get("version", "1"),
    )

    # Delete existing trigger if present
    try:
        client.workflows.triggers.delete(external_id=trigger_external_id)
    except Exception:
        pass

    # Create trigger
    created_trigger = client.workflows.triggers.upsert(trigger, trigger_credentials)
    print(f"Created trigger: {created_trigger.external_id}")
    print(f"  Schedule: {cron_expression}")

    return trigger_external_id, True


def main():
    parser = argparse.ArgumentParser(description="Deploy demo data generator function and workflow")
    parser.add_argument("--dry-run", action="store_true", help="Show what would be done without actually doing it")
    parser.add_argument("--force", action="store_true", help="Force function redeployment even if code hasn't changed")
    parser.add_argument("--output", type=str, help="Output JSON file for results")
    args = parser.parse_args()

    repo_root = get_repo_root()
    print(f"Repository root: {repo_root}")

    # Load configuration
    print("\nLoading configuration...")
    settings = load_settings(repo_root)

    # Check if demo data is enabled
    demo_settings = settings.get("demo_data", {})
    if not demo_settings.get("enabled", False):
        print("\nDemo data generation is DISABLED in settings.yaml")
        print("Set demo_data.enabled: true to enable deployment")
        return 0

    print("\nDemo data generation is ENABLED")

    result = {
        "enabled": True,
        "function_deployed": False,
        "workflow_deployed": False,
        "trigger_deployed": False,
        "dry_run": args.dry_run,
    }

    # Connect to CDF
    print("\nConnecting to CDF...")
    if args.dry_run:
        print("  [DRY RUN] Would connect to CDF")
        client = None
    else:
        client = create_client()
        print(f"  Connected to project: {client.config.project}")

    # Deploy function
    print("\n" + "=" * 60)
    print("[1/3] Deploying Function")
    print("=" * 60)

    if client or args.dry_run:
        deployed = deploy_function(client, repo_root, settings, force=args.force, dry_run=args.dry_run)
        result["function_deployed"] = deployed

    # Deploy workflow
    print("\n" + "=" * 60)
    print("[2/3] Deploying Workflow")
    print("=" * 60)

    if client or args.dry_run:
        workflow_id, workflow_deployed = deploy_workflow(client, settings, force=args.force, dry_run=args.dry_run)
        result["workflow_deployed"] = workflow_deployed
        result["workflow_id"] = workflow_id

    # Deploy trigger
    print("\n" + "=" * 60)
    print("[3/3] Deploying Scheduled Trigger")
    print("=" * 60)

    if client or args.dry_run:
        trigger_id, trigger_deployed = deploy_trigger(client, settings, workflow_id, dry_run=args.dry_run)
        result["trigger_deployed"] = trigger_deployed
        result["trigger_id"] = trigger_id

    # Summary
    print("\n" + "=" * 60)
    print("DEPLOYMENT SUMMARY")
    print("=" * 60)
    print(f"Function deployed: {result['function_deployed']}")
    print(f"Workflow deployed: {result['workflow_deployed']}")
    print(f"Trigger deployed: {result['trigger_deployed']}")

    if args.dry_run:
        print("\n[DRY RUN] No changes were made")

    # Write output if requested
    if args.output:
        import json

        with open(args.output, "w") as f:
            json.dump(result, f)
        print(f"\nResults written to: {args.output}")

    return 0


if __name__ == "__main__":
    sys.exit(main())
