import os
import time
import json
import redis
from celery import Celery
from celery.signals import task_prerun, task_postrun, task_failure
from app.logging import setup_logging, get_logger
from opentelemetry import trace

# Setup logging for Celery worker
logger = setup_logging(service_name="devserver.celery.worker", log_level="DEBUG")

logger.info("=" * 60)
logger.info("DevServer Celery Worker Initializing...")
logger.info("=" * 60)
logger.info("Service: devserver.celery.worker")

# Get tracer for worker
tracer = trace.get_tracer("devserver.celery.tracer")

app = Celery(
    "DevServerTasks",
    broker=os.getenv("CELERY_BROKER_URL"),
    backend=os.getenv("CELERY_RESULT_BACKEND"),
)

logger.info(f"Celery broker: {os.getenv('CELERY_BROKER_URL')}")
logger.info(f"Celery backend: {os.getenv('CELERY_RESULT_BACKEND')}")

# Initialize Redis client for worker
logger.info("Initializing Redis connection for worker...")
try:
    redis_client = redis.Redis(host='redis', port=6379, db=0, decode_responses=True)
    redis_client.ping()
    logger.info("Redis connection established successfully")
except Exception as e:
    logger.error(f"Failed to connect to Redis: {e}", exc_info=True)


# Celery signal handlers
@task_prerun.connect
def task_prerun_handler(sender=None, task_id=None, task=None, args=None, kwargs=None, **extra):
    """Log before task execution"""
    task_name = task.name if task else sender.name if sender else "unknown"  # type: ignore
    logger.info(f"⏳ Task starting: {task_name}")
    logger.info(f"   Service: devserver.celery.worker")
    logger.info(f"   Task ID: {task_id}")
    logger.debug(f"   Args: {args}")
    logger.debug(f"   Kwargs: {kwargs}")


@task_postrun.connect
def task_postrun_handler(sender=None, task_id=None, task=None, args=None, kwargs=None, retval=None, **extra):
    """Log after task execution"""
    task_name = task.name if task else sender.name if sender else "unknown"  # type: ignore
    logger.info(f"Task completed: {task_name}")
    logger.info(f"   Service: devserver.celery.worker")
    logger.info(f"   Task ID: {task_id}")
    logger.debug(f"   Result: {retval}")


@task_failure.connect
def task_failure_handler(sender=None, task_id=None, exception=None, args=None, kwargs=None, traceback=None, **extra):
    """Log task failures"""
    task_name = sender.name if sender else "unknown"  # type: ignore
    logger.error(f"Task failed: {task_name}")
    logger.error(f"   Service: devserver.celery.worker")
    logger.error(f"   Task ID: {task_id}")
    logger.error(f"   Exception: {exception}")
    logger.error(f"   Args: {args}")
    logger.error(f"   Traceback: {traceback}")


@app.task
def long_running_task(duration):
    """Simulates a long-running task by sleeping for the specified duration."""
    logger.info(f"Executing long_running_task: duration={duration}s")
    start_time = time.time()
    
    time.sleep(duration)
    
    elapsed = time.time() - start_time
    result = f"Task completed after {duration} seconds (actual: {elapsed:.2f}s)"
    logger.info(f"{result}")
    return result


@app.task(name="devserver.celery.tasks.fibonacci")
def calculate_fibonacci(n):
    """Calculate the nth Fibonacci number (with intentional delay for demo) and cache the result."""
    with tracer.start_as_current_span("celery.task.fibonacci") as span:
        span.set_attribute("service.name", "devserver.celery.worker")
        span.set_attribute("task.name", "fibonacci")
        span.set_attribute("fibonacci.n", n)
        
        cache_key = f"fibonacci:result:{n}"
        logger.info(f"Calculating Fibonacci number: n={n}, cache_key={cache_key}")
        start_time = time.time()
        
        # Check cache first
        try:
            cached_result = redis_client.get(cache_key)
            if cached_result:
                span.set_attribute("cache.hit", True)
                result_value = int(str(cached_result))
                logger.info(f"Cache hit in worker! Fibonacci({n}) = {result_value}")
                elapsed = time.time() - start_time
                return {
                    "service": "devserver.celery.worker",
                    "task": "fibonacci",
                    "n": n,
                    "result": result_value,
                    "computation_time_seconds": elapsed,
                    "source": "cache"
                }
        except (ValueError, TypeError) as e:
            logger.warning(f"Cache value conversion failed in worker: {e}")
        except Exception as e:
            logger.warning(f"Cache check failed in worker: {e}")
        
        span.set_attribute("cache.hit", False)
        
        def fib(num):
            if num <= 1:
                return num
            return fib(num - 1) + fib(num - 2)
        
        # Add a small delay to make the task more visible in traces
        logger.debug("Adding artificial delay for visibility (0.5s)")
        time.sleep(0.5)
        
        logger.debug(f"Starting Fibonacci calculation for n={n}")
        result = fib(n)
        
        elapsed = time.time() - start_time
        span.set_attribute("fibonacci.result", result)
        span.set_attribute("task.duration_seconds", elapsed)
        
        logger.info(f"Fibonacci calculated: n={n}, result={result}, time={elapsed:.2f}s")
        
        # Store result in cache (TTL: 1 hour for Fibonacci results)
        try:
            redis_client.setex(cache_key, 3600, str(result))
            span.set_attribute("cache.stored", True)
            logger.info(f"Stored Fibonacci result in cache: {cache_key} = {result} (TTL: 3600s)")
        except Exception as e:
            logger.error(f"Failed to store Fibonacci result in cache: {e}")
        
        return {
            "service": "devserver.celery.worker",
            "task": "fibonacci",
            "n": n,
            "result": result,
            "message": f"The {n}th Fibonacci number is {result}",
            "computation_time_seconds": elapsed,
            "source": "computation"
        }


@app.task(name="devserver.celery.tasks.process_data")
def process_data_with_cache(data, cache_key):
    """Process data and cache the result in Redis."""
    with tracer.start_as_current_span("celery.task.process_data") as span:
        span.set_attribute("service.name", "devserver.celery.worker")
        span.set_attribute("task.name", "process_data")
        span.set_attribute("cache.key", cache_key)
        span.set_attribute("data.keys_count", len(data))
        
        logger.info(f"Processing data with caching")
        logger.debug(f"Cache key: {cache_key}")
        logger.debug(f"Data keys: {list(data.keys())}")
        start_time = time.time()
        
        # Simulate processing
        logger.debug("Simulating data processing (1s delay)")
        time.sleep(1)
        
        # Process the data (example: count keys and values)
        processed = {
            "service": "devserver.celery.worker",
            "task": "process_data",
            "input": data,
            "num_keys": len(data),
            "total_length": sum(len(str(v)) for v in data.values()),
            "processed_at": time.time()
        }
        
        logger.debug(f"Data processed: num_keys={processed['num_keys']}, total_length={processed['total_length']}")
        
        # Cache the result for 5 minutes
        try:
            redis_client.setex(cache_key, 300, json.dumps(processed))
            span.set_attribute("cache.operation", "success")
            logger.info(f"Result cached successfully: cache_key={cache_key}, ttl=300s")
        except Exception as e:
            span.set_attribute("cache.operation", "failed")
            logger.error(f"Failed to cache result: {e}", exc_info=True)
        
        elapsed = time.time() - start_time
        span.set_attribute("task.duration_seconds", elapsed)
        logger.info(f"Data processing complete: time={elapsed:.2f}s")
        
        return processed