"""
Data Quality Validation Dashboard

Visualizes historic batch validation progress and results.
Displays orchestration status, partition progress, and validation results.

Runs in CDF Streamlit apps - CogniteClient() is pre-initialized.
"""

from datetime import datetime

import pandas as pd
import streamlit as st
from cognite.client import CogniteClient

# ============================================================================
# Configuration
# ============================================================================
# Space where data quality containers/views are deployed
# Default: "dataQuality" (cog-ai)
# For abp-dev: "sp_data_quality_demo"
DATA_QUALITY_SPACE = "dataQuality"

# Initialize client (automatically available in CDF Streamlit apps)
client = CogniteClient()

st.set_page_config(page_title="Data Quality Validation Dashboard", page_icon="✅", layout="wide")

st.title("📊 Data Quality Validation Dashboard")
st.markdown("Monitor historic batch validation progress and results")

# Sidebar
st.sidebar.header("Actions")
if st.sidebar.button("🔄 Refresh Data", use_container_width=True):
    st.rerun()

st.sidebar.divider()

st.sidebar.header("Filters")
status_filter = st.sidebar.multiselect(
    "Orchestration Status",
    options=["monitoring", "in_progress", "completed", "failed"],
    default=["monitoring", "in_progress", "completed", "failed"],
)

st.sidebar.divider()

st.sidebar.header("About")
st.sidebar.markdown("""
**Data Sources:**
- OrchestrationState (orchestration runs)
- FunctionValidationState (active partitions)
- Function Call Results (completed partitions)

**Key Features:**
- ✅ Multi-orchestration tracking
- ✅ Real-time partition progress
- ✅ Cumulative instance counts
- ✅ Retry and failure monitoring
""")


def query_orchestration_states():
    """Query all orchestration states from the OrchestrationStateView."""
    try:
        from cognite.client.data_classes.data_modeling.ids import ViewId

        # Query via VIEW using SDK
        view_id = ViewId(DATA_QUALITY_SPACE, "OrchestrationStateView", "1")
        result = client.data_modeling.instances.list(instance_type="node", sources=view_id, limit=100)

        orchestrations = []
        for node in result:
            if not node.properties:
                continue

            # Access properties using ViewId as key (not string)
            props = node.properties.get(view_id)
            if props:
                orchestrations.append(
                    {
                        "external_id": node.external_id,
                        "orchestration_id": props.get("orchestration_id", "?"),
                        "view": props.get("view_external_id", "?"),
                        "status": props.get("status", "?"),
                        "created_time": props.get("created_time", "?"),
                        "last_updated": props.get("last_updated", "?"),
                        "partition_count": props.get("partition_count", 0),
                        "total_instances": props.get("total_instances", 0),
                        "validated_instances": props.get("validated_instances", 0),
                        "total_violations": props.get("total_violations", 0),
                        "completed_partitions": props.get("completed_partitions", {}),
                        "failed_partitions": props.get("failed_partitions", {}),
                        "partition_call_ids": props.get("partition_call_ids", {}),
                        "partition_retry_count": props.get("partition_retry_count", {}),
                    }
                )

        return orchestrations
    except Exception as e:
        st.error(f"Error querying orchestration states: {e}")
        return []


def query_cursor_states(orchestration_id=None):
    """
    Query cursor states from the FunctionValidationStateView.

    Args:
        orchestration_id: Optional filter to get states for specific orchestration

    Returns:
        Dict mapping partition_id to cursor state info
    """
    try:
        from cognite.client.data_classes.data_modeling.ids import ViewId

        # Query via VIEW using SDK
        view_id = ViewId(DATA_QUALITY_SPACE, "FunctionValidationStateView", "1")
        result = client.data_modeling.instances.list(instance_type="node", sources=view_id, limit=1000)

        cursors = {}
        for node in result:
            if not node.properties:
                continue

            # Access properties using ViewId as key (not string)
            props = node.properties.get(view_id)
            if not props:
                continue

            orch_id = props.get("orchestration_id", "")

            # Filter by orchestration_id if specified (Python-side)
            if orchestration_id and orch_id != orchestration_id:
                continue

            partition_id = props.get("partitionId", "?")
            cursors[partition_id] = {
                "orchestration_id": orch_id,
                "processed_count": props.get("processedCount", 0),
                "last_updated": props.get("lastUpdated", 0),
            }

        return cursors
    except Exception as e:
        st.warning(f"Could not load cursor states: {e}")
        return {}


# Main content
st.header("Historic Validation Orchestrations")

orchestrations = query_orchestration_states()

# Apply status filter
if status_filter:
    orchestrations = [o for o in orchestrations if o["status"] in status_filter]

if not orchestrations:
    st.info(
        "No orchestration runs found matching the selected filters. Run historic validation from the notebook to see data here."
    )
else:
    # Summary metrics across all orchestrations
    col1, col2, col3, col4, col5 = st.columns(5)

    total_runs = len(orchestrations)
    active_runs = sum(1 for o in orchestrations if o["status"] in ["monitoring", "in_progress"])
    completed_runs = sum(1 for o in orchestrations if o["status"] == "completed")
    failed_runs = sum(1 for o in orchestrations if o["status"] == "failed")
    total_validated = sum(o["validated_instances"] for o in orchestrations)
    total_violations = sum(o["total_violations"] for o in orchestrations)

    col1.metric("Total Runs", total_runs)
    col2.metric("Active", active_runs, delta=None if active_runs == 0 else "🔄")
    col3.metric("Completed", completed_runs, delta="✅" if completed_runs > 0 else None)
    col4.metric("Total Validated", f"{total_validated:,}")
    col5.metric("Total Violations", f"{total_violations:,}", delta="⚠️" if total_violations > 0 else "✓")

    st.divider()

    # Show each orchestration with expandable partition details
    for orch in orchestrations:
        completed_count = len(orch["completed_partitions"])
        failed_count = len(orch["failed_partitions"])
        total_parts = orch["partition_count"]
        in_progress_count = total_parts - completed_count - failed_count
        progress_pct = (completed_count / total_parts * 100) if total_parts > 0 else 0

        # Status emoji
        status_emoji = {"completed": "✅", "failed": "❌", "in_progress": "🔄", "monitoring": "🔄"}.get(
            orch["status"], "❓"
        )

        # Expander title with key info
        expander_title = f"{status_emoji} **{orch['view']}** - {orch['status']} ({progress_pct:.0f}%) - {orch['validated_instances']:,} validated, {orch['total_violations']:,} violations"

        with st.expander(expander_title, expanded=False):
            # Partition metrics (more important, show first)
            col1, col2, col3, col4 = st.columns(4)
            col1.metric("Total Partitions", total_parts)
            col2.metric("Completed", completed_count)
            col3.metric("Failed", failed_count)
            col4.metric("In Progress", in_progress_count)

            # Progress bar
            progress = completed_count / total_parts if total_parts > 0 else 0
            st.progress(progress, text=f"Overall Progress: {progress * 100:.1f}%")

            # Orchestration details (compact format)
            st.caption(f"**Orchestration ID:** `{orch['orchestration_id']}`")

            created_str = orch["created_time"][:16].replace("T", " ") if isinstance(orch["created_time"], str) else "?"
            updated_str = orch["last_updated"][:16].replace("T", " ") if isinstance(orch["last_updated"], str) else "?"
            st.caption(f"**Created:** {created_str} | **Last Updated:** {updated_str}")

            st.markdown("##### Partition Details")

            # Load cursor states for this specific orchestration (for in-progress partitions)
            cursor_states = query_cursor_states(orchestration_id=orch["orchestration_id"])

            # Partition table
            partition_data = []
            for partition_id in range(orch["partition_count"]):
                partition_str = str(partition_id)

                # Determine status
                if partition_str in orch["completed_partitions"]:
                    status = "✅ Completed"
                elif partition_str in orch["failed_partitions"]:
                    status = "❌ Failed"
                else:
                    status = "🔄 In Progress"

                # Get call ID
                call_id = orch["partition_call_ids"].get(partition_str, "N/A")

                # Get retry count
                retry_count = orch["partition_retry_count"].get(partition_str, 0)

                # Try to get processed count
                processed = 0
                last_updated_str = "-"

                # For in-progress partitions, try FunctionValidationState
                if status == "🔄 In Progress":
                    cursor = cursor_states.get(partition_str, {})
                    processed = cursor.get("processed_count", 0)
                    last_updated_ms = cursor.get("last_updated", 0)

                    if last_updated_ms > 0:
                        last_updated_dt = datetime.fromtimestamp(last_updated_ms / 1000)
                        last_updated_str = last_updated_dt.strftime("%Y-%m-%d %H:%M:%S")

                # For completed/failed partitions, try to get from function call result
                if status in ["✅ Completed", "❌ Failed"] and call_id != "N/A":
                    try:
                        call = client.functions.calls.retrieve(
                            call_id=int(call_id), function_external_id="validate-instances-shacl-partitioned"
                        )
                        if hasattr(call, "get_response"):
                            response = call.get_response()
                        else:
                            response = call.response

                        if response:
                            # Use cumulative_validated if available, otherwise instances_validated
                            processed = response.get("cumulative_validated", response.get("instances_validated", 0))

                            # Get end time from call
                            if hasattr(call, "end_time") and call.end_time:
                                last_updated_str = call.end_time.strftime("%Y-%m-%d %H:%M:%S")
                    except Exception:
                        # If we can't fetch the call, just show 0
                        pass

                partition_data.append(
                    {
                        "Partition": partition_id,
                        "Status": status,
                        "Processed": f"{processed:,}" if processed > 0 else "-",
                        "Retries": retry_count,
                        "Last Updated": last_updated_str,
                    }
                )

            df_partitions = pd.DataFrame(partition_data)

            # Style based on status
            def highlight_partition_status(row):
                if "✅" in row["Status"]:
                    return ["background-color: #d4edda"] * len(row)
                elif "❌" in row["Status"]:
                    return ["background-color: #f8d7da"] * len(row)
                elif "🔄" in row["Status"]:
                    return ["background-color: #fff3cd"] * len(row)
                return [""] * len(row)

            st.dataframe(
                df_partitions.style.apply(highlight_partition_status, axis=1), use_container_width=True, hide_index=True
            )
