"""Integration between DAG execution and the job queue.

This module provides:
- DAGQueue: Submit DAGs as jobs to the queue
- DAGJobHandler: Execute DAG jobs with queue integration
- Utilities for converting between DAG jobs and queue jobs
"""

from __future__ import annotations

import logging
from typing import Any, Callable

from oclawma.dag import (
    DAG,
    DAGError,
    DAGExecutor,
    DAGJob,
    DAGResult,
    DAGStatus,
)
from oclawma.queue.models import Job, JobPriority
from oclawma.queue.queue import JobQueue

logger = logging.getLogger(__name__)


class DAGQueueError(Exception):
    """Exception for DAG queue operations."""

    pass


class DAGQueue:
    """Integration between DAG workflows and the job queue.

    This class allows submitting DAGs as jobs to the queue, where the DAG
    is executed with proper dependency management and status tracking.
    """

    def __init__(
        self,
        queue: JobQueue,
        max_workers: int = 4,
        continue_on_error: bool = False,
    ):
        """Initialize DAG queue integration.

        Args:
            queue: The JobQueue instance to use
            max_workers: Maximum parallel workers for DAG execution
            continue_on_error: Whether to continue on job failure
        """
        self.queue = queue
        self.max_workers = max_workers
        self.continue_on_error = continue_on_error

    def submit_dag(
        self,
        dag: DAG,
        handler: Callable[[DAGJob], Any],
        priority: JobPriority = JobPriority.NORMAL,
        scheduled_at: Any | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> Job:
        """Submit a DAG as a job to the queue.

        The DAG is serialized and stored as a job. When dequeued, the entire
        DAG is executed with proper dependency management.

        Args:
            dag: The DAG to execute
            handler: Function to execute each DAG job
            priority: Job priority in the queue
            scheduled_at: Optional scheduled execution time
            metadata: Additional metadata to store

        Returns:
            The created queue job

        Raises:
            DAGQueueError: If DAG validation fails
        """
        try:
            dag.validate()
        except DAGError as e:
            raise DAGQueueError(f"DAG validation failed: {e}") from e

        # Serialize the DAG
        dag_data = self._serialize_dag(dag)

        # Store handler reference (will be resolved at execution time)
        payload = {
            "__dag_execution__": True,
            "dag_data": dag_data,
            "dag_id": dag.dag_id,
            "job_count": dag.job_count,
            "metadata": metadata or {},
        }

        job = self.queue.enqueue(
            payload=payload,
            priority=priority,
            scheduled_at=scheduled_at,
            job_type="dag_workflow",
        )

        logger.info(f"Submitted DAG '{dag.dag_id}' as job {job.id}")
        return job

    def execute_dag_job(
        self,
        job: Job,
        handler: Callable[[DAGJob], Any],
    ) -> DAGResult:
        """Execute a DAG from a queue job.

        Args:
            job: The queue job containing the DAG
            handler: Function to execute each DAG job

        Returns:
            DAGResult with execution results

        Raises:
            DAGQueueError: If job is not a DAG job
        """
        if not self.is_dag_job(job):
            raise DAGQueueError(f"Job {job.id} is not a DAG job")

        dag_data = job.payload["dag_data"]
        dag = self._deserialize_dag(dag_data)

        logger.info(f"Executing DAG '{dag.dag_id}' from job {job.id}")

        executor = DAGExecutor(
            max_workers=self.max_workers,
            continue_on_error=self.continue_on_error,
        )

        return executor.execute_dag(dag, handler)

    def is_dag_job(self, job: Job) -> bool:
        """Check if a job is a DAG execution job."""
        return job.payload.get("__dag_execution__") is True

    def _serialize_dag(self, dag: DAG) -> dict[str, Any]:
        """Serialize a DAG to a dictionary."""
        return {
            "dag_id": dag.dag_id,
            "jobs": [
                {
                    "id": job.id,
                    "payload": job.payload,
                    "priority": job.priority.value,
                    "depends_on": job.depends_on,
                    "max_retries": job.max_retries,
                    "job_type": job.job_type,
                    "metadata": job.metadata,
                }
                for job in dag.jobs.values()
            ],
        }

    def _deserialize_dag(self, data: dict[str, Any]) -> DAG:
        """Deserialize a DAG from a dictionary."""
        dag = DAG(dag_id=data["dag_id"])

        # First pass: create all jobs
        for job_data in data["jobs"]:
            dag.add_job(
                job_id=job_data["id"],
                payload=job_data["payload"],
                priority=JobPriority(job_data["priority"]),
                depends_on=job_data.get("depends_on", []),
                max_retries=job_data.get("max_retries", 3),
                job_type=job_data.get("job_type", "default"),
                metadata=job_data.get("metadata", {}),
            )

        return dag


def create_dag_from_jobs(
    jobs: list[Job],
    dag_id: str | None = None,
) -> DAG:
    """Create a DAG from existing queue jobs that have dependencies.

    This creates a DAG representation of jobs with their dependencies
    from the queue, useful for visualization and analysis.

    Args:
        jobs: List of jobs from the queue
        dag_id: Optional DAG ID

    Returns:
        A DAG representing the job dependencies
    """
    dag = DAG(dag_id=dag_id or "queue_dag")

    for job in jobs:
        if job.id is None:
            continue

        # Convert depends_on to list of strings if needed
        depends_on_list: list[str] = [str(d) for d in job.depends_on] if job.depends_on else []

        dag.add_job(
            job_id=str(job.id),
            payload=job.payload,
            priority=job.priority,
            depends_on=depends_on_list,
            max_retries=job.max_retries,
            job_type=job.job_type,
            metadata={
                "status": job.status.value,
                "retry_count": job.retry_count,
                "error": job.error,
            },
        )

    return dag


def execute_dag_with_queue(
    dag: DAG,
    queue: JobQueue,
    handler: Callable[[DAGJob], Any],
    max_workers: int = 4,
    continue_on_error: bool = False,
) -> DAGResult:
    """Execute a DAG using the job queue for each individual job.

    This submits each DAG job as a separate queue job with dependencies,
    allowing the queue to manage execution and retries.

    Args:
        dag: The DAG to execute
        queue: The JobQueue to use
        handler: Function to process completed jobs
        max_workers: Maximum parallel workers
        continue_on_error: Whether to continue on failure

    Returns:
        DAGResult with execution results
    """
    dag.validate()

    # Submit jobs to queue in topological order
    job_mapping: dict[str, int] = {}  # dag_job_id -> queue_job_id

    for level in dag.topological_sort():
        for dag_job_id in level:
            dag_job = dag.get_job(dag_job_id)

            # Map DAG dependencies to queue job IDs
            queue_depends_on = None
            if dag_job.depends_on:
                queue_depends_on = [
                    job_mapping[dep_id] for dep_id in dag_job.depends_on if dep_id in job_mapping
                ]

            # Create queue job
            queue_job = queue.enqueue(
                payload={
                    "__dag_job__": True,
                    "dag_id": dag.dag_id,
                    "dag_job_id": dag_job.id,
                    "dag_job_payload": dag_job.payload,
                },
                priority=dag_job.priority,
                depends_on=queue_depends_on,
                max_retries=dag_job.max_retries,
                job_type=dag_job.job_type,
            )

            job_mapping[dag_job_id] = queue_job.id

    # TODO: Monitor queue jobs and collect results
    # For now, return a pending result
    return DAGResult(
        dag_id=dag.dag_id,
        status=DAGStatus.PENDING,
        job_results={},
        error="Jobs submitted to queue - monitoring not yet implemented",
    )


def dag_to_dict(dag: DAG) -> dict[str, Any]:
    """Convert a DAG to a dictionary representation.

    This is useful for serialization, APIs, and visualization.

    Args:
        dag: The DAG to convert

    Returns:
        Dictionary representation of the DAG
    """
    return {
        "dag_id": dag.dag_id,
        "job_count": dag.job_count,
        "jobs": [
            {
                "id": job.id,
                "payload": job.payload,
                "priority": {
                    "name": job.priority.name,
                    "value": job.priority.value,
                },
                "depends_on": job.depends_on,
                "max_retries": job.max_retries,
                "job_type": job.job_type,
                "metadata": job.metadata,
                "dependents": list(dag.get_dependents(job.id)),
            }
            for job in dag.jobs.values()
        ],
        "levels": dag.topological_sort(),
    }


def dict_to_dag(data: dict[str, Any]) -> DAG:
    """Create a DAG from a dictionary representation.

    Args:
        data: Dictionary representation of a DAG

    Returns:
        Reconstructed DAG
    """
    dag = DAG(dag_id=data.get("dag_id"))

    for job_data in data["jobs"]:
        priority_value = job_data["priority"]
        if isinstance(priority_value, dict):
            priority_value = priority_value["value"]

        dag.add_job(
            job_id=job_data["id"],
            payload=job_data.get("payload", {}),
            priority=JobPriority(priority_value),
            depends_on=job_data.get("depends_on", []),
            max_retries=job_data.get("max_retries", 3),
            job_type=job_data.get("job_type", "default"),
            metadata=job_data.get("metadata", {}),
        )

    return dag
