from typing import Any, Dict, List
import json

from loguru import logger
from app.core.config import settings

try:
    import redis
except Exception:  # pragma: no cover - optional dependency
    redis = None


class InMemoryTaskStore:
    def __init__(self) -> None:
        self._tasks: Dict[str, Dict[str, Any]] = {}

    def create(self, task_id: str, user_id: str) -> Dict[str, Any]:
        task = {
            "status": "pending",
            "task_id": task_id,
            "progress": 0,
            "user_id": user_id,
        }
        self._tasks[task_id] = task
        return task

    def get(self, task_id: str) -> Dict[str, Any]:
        return self._tasks.get(task_id, {"status": "not_found"})

    def list(self) -> List[Dict[str, Any]]:
        return list(self._tasks.values())

    def update(self, task_id: str, updates: Dict[str, Any]) -> Dict[str, Any]:
        task = self._tasks.get(task_id)
        if not task:
            task = {"status": "pending", "task_id": task_id}
            self._tasks[task_id] = task
        task.update(updates)
        return task


class RedisTaskStore:
    def __init__(self, client: "redis.Redis") -> None:
        self._client = client
        self._index_key = "vigent:tasks:index"

    def _key(self, task_id: str) -> str:
        return f"vigent:tasks:{task_id}"

    def create(self, task_id: str, user_id: str) -> Dict[str, Any]:
        task = {
            "status": "pending",
            "task_id": task_id,
            "progress": 0,
            "user_id": user_id,
        }
        self._client.set(self._key(task_id), json.dumps(task, ensure_ascii=False), ex=86400)
        self._client.sadd(self._index_key, task_id)
        return task

    def get(self, task_id: str) -> Dict[str, Any]:
        raw = self._client.get(self._key(task_id))
        if not raw:
            return {"status": "not_found"}
        return json.loads(raw)

    def list(self) -> List[Dict[str, Any]]:
        task_ids = list(self._client.smembers(self._index_key) or [])
        if not task_ids:
            return []
        keys = [self._key(task_id) for task_id in task_ids]
        raw_items = self._client.mget(keys)
        tasks = []
        expired = []
        for task_id, raw in zip(task_ids, raw_items):
            if raw is None:
                expired.append(task_id)
                continue
            try:
                tasks.append(json.loads(raw))
            except Exception:
                continue
        if expired:
            self._client.srem(self._index_key, *expired)
        return tasks

    def update(self, task_id: str, updates: Dict[str, Any]) -> Dict[str, Any]:
        task = self.get(task_id)
        if task.get("status") == "not_found":
            task = {"status": "pending", "task_id": task_id}
        task.update(updates)
        ttl = 7200 if task.get("status") in ("completed", "failed") else 86400
        self._client.set(self._key(task_id), json.dumps(task, ensure_ascii=False), ex=ttl)
        self._client.sadd(self._index_key, task_id)
        return task


def _build_task_store():
    if redis is None:
        logger.warning("Redis not available, using in-memory task store")
        return InMemoryTaskStore()
    try:
        client = redis.Redis.from_url(settings.REDIS_URL, decode_responses=True)
        client.ping()
        logger.info("Using Redis task store")
        return RedisTaskStore(client)
    except Exception as e:
        logger.warning(f"Redis connection failed, using in-memory task store: {e}")
        return InMemoryTaskStore()


task_store = _build_task_store()


def create_task(task_id: str, user_id: str) -> Dict[str, Any]:
    return task_store.create(task_id, user_id)


def get_task(task_id: str) -> Dict[str, Any]:
    return task_store.get(task_id)


def list_tasks() -> List[Dict[str, Any]]:
    return task_store.list()
