from __future__ import annotations

from typing import TYPE_CHECKING

from cognite.client.data_classes.data_modeling import (
    ContainerApply,
    ContainerProperty,
    DirectRelation,
    Int32,
    Int64,
    Json,
    SpaceApply,
    Text,
    Timestamp,
)

if TYPE_CHECKING:
    from cognite.client import CogniteClient


def _ensure_space_exists(client: CogniteClient, space: str) -> None:
    """Ensure a space exists, creating it if necessary.

    Args:
        client: CogniteClient instance
        space: Space identifier
    """
    try:
        existing_space = client.data_modeling.spaces.retrieve(space)
        if existing_space:
            print(f"  ✓ Space '{space}' already exists")
        else:
            print(f"  • Space '{space}' does not exist, creating...")
            client.data_modeling.spaces.apply(SpaceApply(space=space, name=space))
            print(f"  ✓ Space '{space}' created successfully")
    except Exception:
        # Space doesn't exist, try to create it
        try:
            print(f"  • Space '{space}' not found, creating...")
            client.data_modeling.spaces.apply(SpaceApply(space=space, name=space))
            print(f"  ✓ Space '{space}' created successfully")
        except Exception as create_error:
            print(f"  ✗ Failed to create space '{space}': {create_error}")


def _safe_apply_container(client: CogniteClient, container: ContainerApply) -> None:
    """Apply a container only if it doesn't already exist.

    This avoids permission errors when trying to update existing containers.
    Logs the current status before attempting any operations.
    """
    container_id = f"{container.space}/{container.external_id}"

    try:
        existing = client.data_modeling.containers.retrieve((container.space, container.external_id))
        if existing:
            print(f"    ✓ Container {container_id} already exists, skipping creation")
            return
        else:
            print(f"    • Container {container_id} does not exist, creating...")
    except Exception:
        # Container doesn't exist or error checking, proceed with creation
        print(f"    • Container {container_id} not found, creating...")

    try:
        client.data_modeling.containers.apply(container)
        print(f"    ✓ Container {container_id} created successfully")
    except Exception as e:
        print(f"    ✗ Failed to create container {container_id}: {e}")


def _ensure_orchestration_state_container(client: CogniteClient, space: str, container_id: str) -> None:
    """Ensure the OrchestrationState container exists.

    Creates a container for tracking state of historic validation orchestrations,
    including partition management, progress tracking, and status monitoring.

    Args:
        client: CogniteClient instance
        space: Space for the container (typically "dataQuality")
        container_id: Container external ID (typically "OrchestrationState")
    """
    # Define properties
    properties = {
        "orchestration_id": ContainerProperty(type=Text(), description="Unique ID for this orchestration job"),
        "view_external_id": ContainerProperty(type=Text(), description="View being validated"),
        "view_space": ContainerProperty(type=Text(), description="Space of the view"),
        "instance_space": ContainerProperty(type=Text(), description="Instance space being validated"),
        "rule_set_id": ContainerProperty(type=Text(), description="Rule set ID for validation results"),
        "workflow_external_id": ContainerProperty(
            type=Text(), description="Scheduled workflow managing this orchestration"
        ),
        "partition_count": ContainerProperty(type=Int32(), description="Total number of partitions"),
        "partition_field": ContainerProperty(type=Text(), description="Field used for partitioning"),
        "partition_ranges": ContainerProperty(
            type=Json(), description="Object mapping partition IDs to range configurations"
        ),
        "triggered_partitions": ContainerProperty(type=Json(), description="Object mapping partition IDs to boolean"),
        "completed_partitions": ContainerProperty(type=Json(), description="Object mapping partition IDs to boolean"),
        "failed_partitions": ContainerProperty(type=Json(), description="Object mapping partition IDs to boolean"),
        "partition_call_ids": ContainerProperty(
            type=Json(), description="Map of partition_id to latest function call_id"
        ),
        "partition_retry_count": ContainerProperty(
            type=Json(), description="Map of partition_id to current retry count"
        ),
        "total_instances": ContainerProperty(type=Int64(), description="Total instances discovered"),
        "validated_instances": ContainerProperty(type=Int64(), description="Total instances validated so far"),
        "total_violations": ContainerProperty(type=Int64(), description="Total violations found so far"),
        "status": ContainerProperty(
            type=Text(), description="Current status (initial | monitoring | completed | failed)"
        ),
        "created_time": ContainerProperty(type=Timestamp(), description="When orchestration was created"),
        "last_updated": ContainerProperty(type=Timestamp(), description="Last time state was updated"),
        "max_retries": ContainerProperty(type=Int32(), description="Maximum retries per partition"),
        "sync_trigger_external_id": ContainerProperty(
            type=Text(), description="Sync trigger created for ongoing validation"
        ),
        "historic_cutoff_date": ContainerProperty(
            type=Text(), description="Maximum date from historic partitions (ISO format)"
        ),
        "timestamp_collection_count": ContainerProperty(
            type=Int64(), description="Number of timestamps collected so far"
        ),
        "timestamp_collection_phase": ContainerProperty(
            type=Text(), description="Collection phase: 'collecting' | 'complete' | null"
        ),
        "timestamp_collection_started": ContainerProperty(
            type=Timestamp(), description="When timestamp collection started (for total time tracking)"
        ),
        "timestamps_file": ContainerProperty(
            type=DirectRelation(),
            description=(
                "CogniteFile node (dataQuality space) holding accumulated timestamps + cursor for resumable collection"
            ),
        ),
    }

    container = ContainerApply(
        space=space,
        external_id=container_id,
        name="Orchestration State",
        properties=properties,
        description="Tracks state for self-managing historic validation orchestrations",
        used_for="node",
    )

    _safe_apply_container(client, container)


def _ensure_function_validation_state_container(client: CogniteClient, space: str, container_id: str) -> None:
    """Ensure the FunctionValidationState container exists.

    Creates a container for storing cursor state for partitioned validation functions,
    enabling resumption after timeouts or failures.

    Args:
        client: CogniteClient instance
        space: Space for the container (typically "dataQuality")
        container_id: Container external ID (typically "FunctionValidationState")
    """
    # Define properties
    properties = {
        "orchestration_id": ContainerProperty(type=Text(), description="Orchestration run this partition belongs to"),
        "partitionId": ContainerProperty(type=Text(), description="Unique partition ID"),
        "syncCursor": ContainerProperty(
            type=Text(), description="DMS sync cursor for resuming from last processed position"
        ),
        "processedCount": ContainerProperty(
            type=Int64(), description="Total number of instances processed so far in this partition"
        ),
        "lastUpdated": ContainerProperty(
            type=Int64(), description="Timestamp (milliseconds since epoch) of last cursor save"
        ),
    }

    container = ContainerApply(
        space=space,
        external_id=container_id,
        name="Function Validation State",
        properties=properties,
        description=(
            "Stores cursor state for partitioned validation functions to enable resumption after timeouts or failures"
        ),
        used_for="node",
    )

    _safe_apply_container(client, container)


def _ensure_raw_validation_state_container(client: CogniteClient, space: str, container_id: str) -> None:
    """Ensure the RawValidationState container exists.

    Creates a container for tracking incremental validation state of RAW tables,
    including last processed timestamp, row counts, and validation status.

    Args:
        client: CogniteClient instance
        space: Space for the container (typically "dataQuality")
        container_id: Container external ID (typically "RawValidationState")
    """
    # Define properties
    properties = {
        "db_name": ContainerProperty(type=Text(), description="RAW database name"),
        "table_name": ContainerProperty(type=Text(), description="RAW table name"),
        "last_processed_timestamp": ContainerProperty(
            type=Int64(), description="Last updated time of most recently validated row (milliseconds)"
        ),
        "last_run_time": ContainerProperty(type=Timestamp(), description="Last validation run timestamp"),
        "rows_processed": ContainerProperty(type=Int64(), description="Number of rows processed in last run"),
        "last_validation_conforms": ContainerProperty(
            type=Text(), description="Whether last validation passed (true/false)"
        ),
        "last_violation_count": ContainerProperty(type=Int64(), description="Number of violations in last run"),
    }

    container = ContainerApply(
        space=space,
        external_id=container_id,
        name="Raw Validation State",
        properties=properties,
        description="Tracks incremental validation state for RAW tables",
        used_for="node",
    )

    _safe_apply_container(client, container)
