# services/job_service.py - Enhanced with atomic operations
import asyncio
from contextlib import asynccontextmanager
from typing import Dict, Any, Optional
from uuid import UUID
import logging

from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.exc import SQLAlchemyError
import redis.asyncio as aioredis
import redis
import rq
from ..settings import RQ_QUEUE_NAME

from ..models import Job, JobStatus, JobType, Project
from ..storage import StorageBackend

log = logging.getLogger(__name__)


def _get_sync_redis():
    """Get a synchronous Redis connection for RQ."""
    import os

    redis_url = os.environ.get("REDIS_URL", "redis://127.0.0.1:6379/0")
    return redis.Redis.from_url(
        redis_url,
        decode_responses=False,
        socket_connect_timeout=10,
        socket_timeout=30,
        retry_on_timeout=True,
    )


class JobCreationError(Exception):
    """Raised when job creation fails"""

    pass


class JobService:
    """Orchestrates job creation with transactional guarantees"""

    def __init__(
        self,
        db_session: AsyncSession,
        redis_client: aioredis.Redis,
        storage: StorageBackend,
        ml_detector,
    ):
        self.db = db_session
        self.redis = redis_client
        self.storage = storage
        self.ml_detector = ml_detector
        # Create RQ queue with SYNCHRONOUS Redis connection (RQ doesn't support async)
        self._sync_redis = _get_sync_redis()
        self.queue = rq.Queue(RQ_QUEUE_NAME, connection=self._sync_redis)

    def _log_queue_state(self, label: str, job_id: Optional[str] = None) -> None:
        """Emit queue length and connection details for debugging."""
        try:
            queue_key = self.queue.key
            length = self._sync_redis.llen(queue_key)
            conn = self._sync_redis.connection_pool.connection_kwargs
            host = conn.get("host")
            port = conn.get("port")
            db = conn.get("db")
            log.info(
                "RQ queue state (%s): redis=%s:%s/%s key=%s len=%s job=%s",
                label,
                host,
                port,
                db,
                queue_key,
                length,
                job_id,
            )
        except Exception as exc:
            log.warning("Failed to log RQ queue state (%s): %s", label, exc)

    def _ensure_job_queued(self, job_id: str) -> None:
        """Ensure the job id is present in the Redis queue list."""
        queue_key = self.queue.key
        job_bytes = job_id.encode("utf-8")
        try:
            try:
                pos = self._sync_redis.lpos(queue_key, job_id)
            except Exception:
                pos = None
            if pos is None:
                try:
                    tail = self._sync_redis.lrange(queue_key, -200, -1)
                except Exception:
                    tail = []
                if tail and isinstance(tail[0], bytes):
                    found = job_bytes in tail
                else:
                    found = str(job_id) in [str(v) for v in tail]
                if not found:
                    self._sync_redis.rpush(queue_key, job_id)
                    log.warning(
                        "RQ enqueue fallback: pushed job %s to %s", job_id, queue_key
                    )
        except Exception as exc:
            log.warning("Failed to verify RQ enqueue for job %s: %s", job_id, exc)

    async def create_and_enqueue_job(
        self, project_id: UUID, job_type: JobType, config: Dict[str, Any]
    ) -> Job:
        """
        Atomically create job in DB and enqueue to Redis.
        Uses two-phase commit pattern for reliability.
        """
        job = None
        enqueued = False

        try:
            # Phase 1: Create job in database (but don't commit yet)
            job = Job(
                project_id=project_id,
                status=JobStatus.QUEUED,
                job_type=job_type,
                config=config,
            )

            self.db.add(job)
            await self.db.flush()  # Get the ID without committing

            # Phase 2: Try to enqueue to Redis using RQ's proper API
            try:
                # Import the worker function
                from ..worker import run_matching_job

                # Use RQ's native enqueue (synchronous, but quick)
                rq_job = self.queue.enqueue(
                    run_matching_job,
                    kwargs={"job_id": str(job.id)},
                    job_id=str(job.id),
                    job_timeout="1h",
                )
                enqueued = True
                log.info(f"Enqueued job {job.id} to RQ, rq_status={rq_job.get_status()}")
                self._log_queue_state("after_enqueue", str(job.id))
                self._ensure_job_queued(str(job.id))
                self._log_queue_state("after_ensure", str(job.id))

                # Also set tracking key via async redis
                await self.redis.set(
                    f"job:{job.id}:queued",
                    "1",
                    ex=86400,  # Expire after 24 hours
                )

            except Exception as e:
                log.error(f"Failed to enqueue job {job.id}: {e}")
                raise JobCreationError(f"Failed to enqueue job: {e}")

            # Phase 3: If enqueuing succeeded, commit the database transaction
            await self.db.commit()
            await self.db.refresh(job)

            log.info(f"Successfully created and enqueued job {job.id}")
            return job

        except SQLAlchemyError as e:
            # Database error - rollback
            log.error(f"Database error creating job: {e}")
            await self.db.rollback()
            raise JobCreationError(f"Database error: {e}")

        except Exception as e:
            # Any other error - ensure cleanup
            log.error(f"Unexpected error creating job: {e}")
            await self.db.rollback()

            # If we enqueued but couldn't commit to DB, try to remove from queue
            if enqueued and job:
                try:
                    await self._remove_from_queue(job.id)
                except Exception as cleanup_error:
                    log.error(f"Failed to cleanup queue after error: {cleanup_error}")

            raise JobCreationError(f"Failed to create job: {e}")

    def _serialize_job_for_rq(self, job_id: str) -> str:
        """Serialize job for RQ protocol"""
        import pickle
        import base64

        # RQ expects a specific format
        job_dict = {
            "func": "fmatch.saas.worker.run_matching_job",
            "args": (),
            "kwargs": {"job_id": job_id},
            "timeout": "1h",
            "id": job_id,
        }

        return base64.b64encode(pickle.dumps(job_dict)).decode("utf-8")

    async def _remove_from_queue(self, job_id: UUID):
        """Emergency cleanup - remove job from queue"""
        try:
            # Remove from RQ queue
            await self.redis.lrem("rq:queue:default", 0, str(job_id))
            # Remove tracking key
            await self.redis.delete(f"job:{job_id}:queued")
        except Exception as e:
            log.error(f"Failed to remove job {job_id} from queue: {e}")

    async def estimate_duration(
        self, source_url: str, reference_url: Optional[str], config: Dict[str, Any]
    ) -> int:
        """Estimate job duration based on file sizes and config"""
        # Simple estimation based on file sizes
        # In production, this could use ML models based on historical data
        try:
            source_size = await self.storage.get_file_size(source_url)
            ref_size = (
                await self.storage.get_file_size(reference_url) if reference_url else 0
            )

            # Base estimate: 1 second per 1000 rows (assuming ~100 bytes per row)
            estimated_rows = (source_size + ref_size) / 100
            base_duration = estimated_rows / 1000

            # Adjust based on algorithm complexity
            algorithm = config.get("algorithm", "WRatio")
            complexity_multiplier = {
                "Exact": 0.5,
                "WRatio": 1.0,
                "TokenSetRatio": 1.5,
                "JaroWinkler": 2.0,
            }.get(algorithm, 1.0)

            # Add overhead for blocking, preprocessing, etc.
            duration = int(base_duration * complexity_multiplier + 30)  # 30s minimum

            return min(duration, 3600)  # Cap at 1 hour

        except Exception as e:
            log.warning(f"Could not estimate duration: {e}")
            return 300  # Default 5 minutes

    async def get_or_create_default_project(self, account_id: UUID) -> Project:
        """Get or create a default project for the account"""
        # Check for existing default project
        from sqlalchemy import select

        stmt = select(Project).where(
            Project.account_id == account_id, Project.name == "Default Project"
        )
        result = await self.db.execute(stmt)
        project = result.scalar_one_or_none()

        if not project:
            # Create default project
            project = Project(
                account_id=account_id,
                name="Default Project",
                description="Automatically created default project",
            )
            self.db.add(project)
            await self.db.flush()

        return project

    async def auto_detect_configuration(
        self,
        source_url: str,
        reference_url: Optional[str],
        mode: str,
        overrides: Dict[str, Any],
    ) -> Dict[str, Any]:
        """Auto-detect optimal configuration using ML"""
        try:
            # Read sample data from storage
            source_sample = await self.storage.read_sample(source_url, max_rows=1000)
            ref_sample = None
            if reference_url:
                ref_sample = await self.storage.read_sample(
                    reference_url, max_rows=1000
                )

            # Use ML detector if available
            if self.ml_detector:
                config = await self.ml_detector.detect_configuration(
                    source_df=source_sample, reference_df=ref_sample, mode=mode
                )
            else:
                # Fallback to basic configuration
                config = {
                    "algorithm": "WRatio",
                    "threshold": 80,
                    "blocking_enabled": True,
                }

            # Apply any user overrides
            config.update(overrides)

            return config

        except Exception as e:
            log.error(f"Auto-detection failed: {e}")
            # Return sensible defaults
            return {
                "algorithm": "WRatio",
                "threshold": 80,
                "blocking_enabled": True,
                **overrides,
            }

    @asynccontextmanager
    async def atomic_job_creation(self):
        """
        Context manager for atomic job creation.
        Ensures rollback on any failure.
        """
        try:
            yield self
        except Exception:
            await self.db.rollback()
            raise
