from __future__ import annotations

import copy
import json
import os
import sqlite3
import threading
from collections import deque
from enum import Enum
from typing import Any, Dict, Optional

from grantflow.core.strategies.factory import DonorFactory
from grantflow.swarm.state_contract import normalize_state_contract

RUNTIME_STATE_KEYS = {"strategy", "donor_strategy"}
DEFAULT_SQLITE_BUSY_TIMEOUT_MS = 5000


def _env(name: str, default: str) -> str:
    legacy = name.replace("GRANTFLOW_", "AIDGRAPH_", 1) if name.startswith("GRANTFLOW_") else name
    return os.getenv(name, os.getenv(legacy, default))


def default_sqlite_path() -> str:
    return _env("GRANTFLOW_SQLITE_PATH", "./grantflow_state.db")


def sqlite_busy_timeout_ms() -> int:
    raw = _env("GRANTFLOW_SQLITE_BUSY_TIMEOUT_MS", str(DEFAULT_SQLITE_BUSY_TIMEOUT_MS))
    try:
        value = int(raw)
    except ValueError:
        return DEFAULT_SQLITE_BUSY_TIMEOUT_MS
    return max(0, value)


def storage_mode(name: str, default: str = "inmem") -> str:
    return _env(name, default).strip().lower()


def sanitize_jsonable(value: Any) -> Any:
    if value is None or isinstance(value, (str, int, float, bool)):
        return value
    if isinstance(value, Enum):
        return value.value
    if isinstance(value, dict):
        return {str(k): sanitize_jsonable(v) for k, v in value.items()}
    if isinstance(value, (list, tuple)):
        return [sanitize_jsonable(item) for item in value]
    return str(value)


def _normalized_tenant_token(value: Any) -> str:
    token = str(value or "").strip().lower()
    return token


def _tenant_from_row(namespace: Any, metadata: Any) -> str:
    metadata_dict = metadata if isinstance(metadata, dict) else {}
    from_meta = _normalized_tenant_token(metadata_dict.get("tenant_id") or metadata_dict.get("tenant"))
    if from_meta:
        return from_meta
    raw_namespace = str(namespace or "").strip()
    if "/" in raw_namespace:
        return _normalized_tenant_token(raw_namespace.split("/", 1)[0])
    return ""


def prepare_state_for_storage(state: Any) -> Any:
    if not isinstance(state, dict):
        return sanitize_jsonable(state)

    normalized_state = copy.deepcopy(state)
    normalize_state_contract(normalized_state, emit_legacy_aliases=True)
    stored_state: Dict[str, Any] = {}
    for key, value in normalized_state.items():
        if key in RUNTIME_STATE_KEYS:
            continue
        stored_state[str(key)] = sanitize_jsonable(value)
    return stored_state


def restore_state_from_storage(state: Any) -> Any:
    if not isinstance(state, dict):
        return state

    restored = copy.deepcopy(state)
    donor_id = restored.get("donor_id") or restored.get("donor")
    if donor_id and "donor_strategy" not in restored and "strategy" not in restored:
        try:
            strategy = DonorFactory.get_strategy(str(donor_id))
        except Exception:
            return restored
        restored["donor_strategy"] = strategy
        restored["strategy"] = strategy
    normalize_state_contract(restored, emit_legacy_aliases=True)
    return restored


def _normalize_state_in_payload(payload: Dict[str, Any]) -> Dict[str, Any]:
    out = copy.deepcopy(payload)
    state = out.get("state")
    if isinstance(state, dict):
        normalize_state_contract(state, emit_legacy_aliases=True)
    return out


def prepare_job_payload_for_storage(payload: Dict[str, Any]) -> Dict[str, Any]:
    stored: Dict[str, Any] = {}
    for key, value in payload.items():
        if key == "state":
            stored[str(key)] = prepare_state_for_storage(value)
        else:
            stored[str(key)] = sanitize_jsonable(value)
    return stored


def restore_job_payload_from_storage(payload: Dict[str, Any]) -> Dict[str, Any]:
    restored = copy.deepcopy(payload)
    if "state" in restored:
        restored["state"] = restore_state_from_storage(restored["state"])
    return restored


def storage_json_dumps(value: Any) -> str:
    return json.dumps(value, ensure_ascii=False, separators=(",", ":"))


def storage_json_loads(value: str) -> Any:
    return json.loads(value)


def open_sqlite_connection(db_path: str) -> sqlite3.Connection:
    conn = sqlite3.connect(db_path, timeout=30)
    conn.row_factory = sqlite3.Row
    conn.execute(f"PRAGMA busy_timeout = {sqlite_busy_timeout_ms()}")
    conn.execute("PRAGMA journal_mode = WAL")
    conn.execute("PRAGMA synchronous = NORMAL")
    return conn


def ensure_sqlite_schema_meta(conn: sqlite3.Connection) -> None:
    conn.execute("""
        CREATE TABLE IF NOT EXISTS schema_meta (
          component TEXT PRIMARY KEY,
          version INTEGER NOT NULL,
          updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
        )
        """)


def ensure_sqlite_component_schema(conn: sqlite3.Connection, component: str, target_version: int) -> int:
    ensure_sqlite_schema_meta(conn)
    row = conn.execute("SELECT version FROM schema_meta WHERE component = ?", (component,)).fetchone()
    if row is None:
        conn.execute(
            """
            INSERT INTO schema_meta (component, version, updated_at)
            VALUES (?, ?, CURRENT_TIMESTAMP)
            """,
            (component, target_version),
        )
        return target_version

    current_version = int(row["version"])
    if current_version > target_version:
        raise RuntimeError(
            f"Unsupported newer schema for component '{component}': {current_version} > {target_version}"
        )
    if current_version < target_version:
        # Placeholder for future migrations; current schema is single-step and idempotent.
        conn.execute(
            """
            UPDATE schema_meta
            SET version = ?, updated_at = CURRENT_TIMESTAMP
            WHERE component = ?
            """,
            (target_version, component),
        )
    return target_version


class InMemoryJobStore:
    def __init__(self) -> None:
        self._jobs: Dict[str, Dict[str, Any]] = {}
        self._lock = threading.Lock()

    def set(self, job_id: str, payload: Dict[str, Any]) -> None:
        normalized_payload = _normalize_state_in_payload(payload)
        with self._lock:
            self._jobs[job_id] = copy.deepcopy(normalized_payload)

    def update(self, job_id: str, **patch: Any) -> Dict[str, Any]:
        with self._lock:
            current = copy.deepcopy(self._jobs.get(job_id, {}))
            current.update(patch)
            normalized = _normalize_state_in_payload(current)
            self._jobs[job_id] = copy.deepcopy(normalized)
            return copy.deepcopy(normalized)

    def get(self, job_id: str) -> Optional[Dict[str, Any]]:
        with self._lock:
            payload = self._jobs.get(job_id)
            return copy.deepcopy(payload) if payload is not None else None

    def list(self) -> Dict[str, Dict[str, Any]]:
        with self._lock:
            return {job_id: copy.deepcopy(payload) for job_id, payload in self._jobs.items()}


class InMemoryIngestAuditStore:
    def __init__(self, maxlen: int = 500) -> None:
        self._rows: deque[Dict[str, Any]] = deque(maxlen=max(1, int(maxlen)))
        self._lock = threading.Lock()

    def append(self, row: Dict[str, Any]) -> None:
        with self._lock:
            self._rows.append(copy.deepcopy(sanitize_jsonable(row)))

    def list_recent(
        self,
        donor_id: Optional[str] = None,
        tenant_id: Optional[str] = None,
        limit: int = 50,
    ) -> list[Dict[str, Any]]:
        donor_filter = str(donor_id or "").strip().lower()
        tenant_filter = _normalized_tenant_token(tenant_id)
        with self._lock:
            rows = list(self._rows)
        rows = list(reversed(rows))
        if donor_filter:
            rows = [r for r in rows if str(r.get("donor_id") or "").strip().lower() == donor_filter]
        if tenant_filter:
            rows = [
                r
                for r in rows
                if _tenant_from_row(
                    r.get("namespace"),
                    r.get("metadata") if isinstance(r.get("metadata"), dict) else {},
                )
                == tenant_filter
            ]
        capped = rows[: max(1, min(int(limit or 50), 200))]
        return [copy.deepcopy(r) for r in capped]

    def clear(self) -> None:
        with self._lock:
            self._rows.clear()

    def inventory(self, donor_id: Optional[str] = None, tenant_id: Optional[str] = None) -> list[Dict[str, Any]]:
        donor_filter = str(donor_id or "").strip().lower()
        tenant_filter = _normalized_tenant_token(tenant_id)
        with self._lock:
            rows = list(self._rows)

        grouped: Dict[str, Dict[str, Any]] = {}
        for row in rows:
            row_donor = str(row.get("donor_id") or "").strip()
            if donor_filter and row_donor.lower() != donor_filter:
                continue
            metadata = row.get("metadata") if isinstance(row.get("metadata"), dict) else {}
            row_tenant = _tenant_from_row(row.get("namespace"), metadata)
            if tenant_filter and row_tenant != tenant_filter:
                continue
            doc_family = str((metadata or {}).get("doc_family") or "").strip()
            if not doc_family:
                continue
            key = f"{row_tenant}::{row_donor.lower()}::{doc_family}"
            ts = str(row.get("ts") or "")
            current = grouped.get(key)
            if current is None:
                grouped[key] = {
                    "tenant_id": row_tenant or None,
                    "donor_id": row_donor,
                    "doc_family": doc_family,
                    "count": 1,
                    "latest_ts": ts,
                    "latest_filename": str(row.get("filename") or ""),
                    "latest_event_id": str(row.get("event_id") or ""),
                    "latest_source_type": str((metadata or {}).get("source_type") or ""),
                }
                continue
            current["count"] = int(current.get("count") or 0) + 1
            if ts >= str(current.get("latest_ts") or ""):
                current["latest_ts"] = ts
                current["latest_filename"] = str(row.get("filename") or "")
                current["latest_event_id"] = str(row.get("event_id") or "")
                current["latest_source_type"] = str((metadata or {}).get("source_type") or "")

        return sorted(
            grouped.values(),
            key=lambda item: (
                str(item.get("tenant_id") or ""),
                str(item.get("donor_id") or ""),
                -int(item.get("count") or 0),
                str(item.get("doc_family") or ""),
            ),
        )


class SQLiteJobStore:
    SCHEMA_COMPONENT = "jobs"
    SCHEMA_VERSION = 1

    def __init__(self, db_path: Optional[str] = None) -> None:
        self.db_path = db_path or default_sqlite_path()
        self._write_lock = threading.Lock()
        self._init_db()

    def _connect(self) -> sqlite3.Connection:
        return open_sqlite_connection(self.db_path)

    def _init_db(self) -> None:
        with self._connect() as conn:
            ensure_sqlite_component_schema(conn, self.SCHEMA_COMPONENT, self.SCHEMA_VERSION)
            conn.execute("""
                CREATE TABLE IF NOT EXISTS jobs (
                  job_id TEXT PRIMARY KEY,
                  payload_json TEXT NOT NULL,
                  updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
                )
                """)

    def set(self, job_id: str, payload: Dict[str, Any]) -> None:
        stored_payload = prepare_job_payload_for_storage(payload)
        payload_json = storage_json_dumps(stored_payload)
        with self._write_lock:
            with self._connect() as conn:
                conn.execute(
                    """
                    INSERT INTO jobs (job_id, payload_json, updated_at)
                    VALUES (?, ?, CURRENT_TIMESTAMP)
                    ON CONFLICT(job_id) DO UPDATE SET
                      payload_json=excluded.payload_json,
                      updated_at=CURRENT_TIMESTAMP
                    """,
                    (job_id, payload_json),
                )

    def update(self, job_id: str, **patch: Any) -> Dict[str, Any]:
        with self._write_lock:
            with self._connect() as conn:
                row = conn.execute("SELECT payload_json FROM jobs WHERE job_id = ?", (job_id,)).fetchone()
                current: Dict[str, Any] = storage_json_loads(row["payload_json"]) if row else {}
                merged = dict(current)
                merged.update(prepare_job_payload_for_storage(patch))
                conn.execute(
                    """
                    INSERT INTO jobs (job_id, payload_json, updated_at)
                    VALUES (?, ?, CURRENT_TIMESTAMP)
                    ON CONFLICT(job_id) DO UPDATE SET
                      payload_json=excluded.payload_json,
                      updated_at=CURRENT_TIMESTAMP
                    """,
                    (job_id, storage_json_dumps(merged)),
                )
        return restore_job_payload_from_storage(merged)

    def get(self, job_id: str) -> Optional[Dict[str, Any]]:
        with self._connect() as conn:
            row = conn.execute("SELECT payload_json FROM jobs WHERE job_id = ?", (job_id,)).fetchone()
        if not row:
            return None
        payload = storage_json_loads(row["payload_json"])
        return restore_job_payload_from_storage(payload)

    def list(self) -> Dict[str, Dict[str, Any]]:
        with self._connect() as conn:
            rows = conn.execute("SELECT job_id, payload_json FROM jobs ORDER BY updated_at DESC").fetchall()
        items: Dict[str, Dict[str, Any]] = {}
        for row in rows:
            items[str(row["job_id"])] = restore_job_payload_from_storage(storage_json_loads(row["payload_json"]))
        return items


class SQLiteIngestAuditStore:
    SCHEMA_COMPONENT = "ingest_audit"
    SCHEMA_VERSION = 1

    def __init__(self, db_path: Optional[str] = None) -> None:
        self.db_path = db_path or default_sqlite_path()
        self._write_lock = threading.Lock()
        self._init_db()

    def _connect(self) -> sqlite3.Connection:
        return open_sqlite_connection(self.db_path)

    def _init_db(self) -> None:
        with self._connect() as conn:
            ensure_sqlite_component_schema(conn, self.SCHEMA_COMPONENT, self.SCHEMA_VERSION)
            conn.execute("""
                CREATE TABLE IF NOT EXISTS ingest_audit_events (
                  event_id TEXT PRIMARY KEY,
                  ts TEXT NOT NULL,
                  donor_id TEXT NOT NULL,
                  namespace TEXT NOT NULL,
                  filename TEXT NOT NULL,
                  content_type TEXT NOT NULL,
                  metadata_json TEXT NOT NULL,
                  result_json TEXT NOT NULL,
                  created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
                )
                """)

    def append(self, row: Dict[str, Any]) -> None:
        item = sanitize_jsonable(dict(row or {}))
        with self._write_lock:
            with self._connect() as conn:
                conn.execute(
                    """
                    INSERT INTO ingest_audit_events (
                      event_id, ts, donor_id, namespace, filename, content_type, metadata_json, result_json, created_at
                    ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
                    ON CONFLICT(event_id) DO UPDATE SET
                      ts=excluded.ts,
                      donor_id=excluded.donor_id,
                      namespace=excluded.namespace,
                      filename=excluded.filename,
                      content_type=excluded.content_type,
                      metadata_json=excluded.metadata_json,
                      result_json=excluded.result_json,
                      created_at=CURRENT_TIMESTAMP
                    """,
                    (
                        str(item.get("event_id") or ""),
                        str(item.get("ts") or ""),
                        str(item.get("donor_id") or ""),
                        str(item.get("namespace") or ""),
                        str(item.get("filename") or ""),
                        str(item.get("content_type") or ""),
                        storage_json_dumps(item.get("metadata") or {}),
                        storage_json_dumps(item.get("result") or {}),
                    ),
                )

    def list_recent(
        self,
        donor_id: Optional[str] = None,
        tenant_id: Optional[str] = None,
        limit: int = 50,
    ) -> list[Dict[str, Any]]:
        cap = max(1, min(int(limit or 50), 200))
        donor_filter = str(donor_id or "").strip()
        tenant_filter = _normalized_tenant_token(tenant_id)
        query = (
            "SELECT event_id, ts, donor_id, namespace, filename, content_type, metadata_json, result_json "
            "FROM ingest_audit_events "
        )
        params: list[Any] = []
        if donor_filter:
            query += "WHERE lower(donor_id) = lower(?) "
            params.append(donor_filter)
        query += "ORDER BY ts DESC, created_at DESC LIMIT ?"
        params.append(cap)
        with self._connect() as conn:
            rows = conn.execute(query, tuple(params)).fetchall()

        items: list[Dict[str, Any]] = []
        for row in rows:
            metadata = storage_json_loads(row["metadata_json"])
            metadata = metadata if isinstance(metadata, dict) else {}
            row_tenant = _tenant_from_row(row["namespace"], metadata)
            if tenant_filter and row_tenant != tenant_filter:
                continue
            items.append(
                {
                    "event_id": row["event_id"],
                    "ts": row["ts"],
                    "donor_id": row["donor_id"],
                    "namespace": row["namespace"],
                    "filename": row["filename"],
                    "content_type": row["content_type"],
                    "metadata": metadata,
                    "result": storage_json_loads(row["result_json"]),
                }
            )
        return items

    def clear(self) -> None:
        with self._write_lock:
            with self._connect() as conn:
                conn.execute("DELETE FROM ingest_audit_events")

    def inventory(self, donor_id: Optional[str] = None, tenant_id: Optional[str] = None) -> list[Dict[str, Any]]:
        donor_filter = str(donor_id or "").strip()
        tenant_filter = _normalized_tenant_token(tenant_id)
        query = "SELECT event_id, ts, donor_id, namespace, filename, metadata_json " "FROM ingest_audit_events "
        params: list[Any] = []
        if donor_filter:
            query += "WHERE lower(donor_id) = lower(?) "
            params.append(donor_filter)
        query += "ORDER BY ts DESC, created_at DESC"
        with self._connect() as conn:
            rows = conn.execute(query, tuple(params)).fetchall()

        grouped: Dict[str, Dict[str, Any]] = {}
        for row in rows:
            row_donor = str(row["donor_id"] or "").strip()
            metadata = storage_json_loads(row["metadata_json"])
            metadata = metadata if isinstance(metadata, dict) else {}
            row_tenant = _tenant_from_row(row["namespace"], metadata)
            if tenant_filter and row_tenant != tenant_filter:
                continue
            doc_family = str(metadata.get("doc_family") or "").strip()
            if not doc_family:
                continue
            key = f"{row_tenant}::{row_donor.lower()}::{doc_family}"
            ts = str(row["ts"] or "")
            current = grouped.get(key)
            if current is None:
                grouped[key] = {
                    "tenant_id": row_tenant or None,
                    "donor_id": row_donor,
                    "doc_family": doc_family,
                    "count": 1,
                    "latest_ts": ts,
                    "latest_filename": str(row["filename"] or ""),
                    "latest_event_id": str(row["event_id"] or ""),
                    "latest_source_type": str(metadata.get("source_type") or ""),
                }
                continue
            current["count"] = int(current.get("count") or 0) + 1

        return sorted(
            grouped.values(),
            key=lambda item: (
                str(item.get("tenant_id") or ""),
                str(item.get("donor_id") or ""),
                -int(item.get("count") or 0),
                str(item.get("doc_family") or ""),
            ),
        )


_REDIS_JOB_KEY_PREFIX = "grantflow:job:"
_REDIS_JOB_INDEX_KEY = "grantflow:job_index"
_REDIS_JOB_TTL_SECONDS = 90 * 24 * 3600  # 90 days


def _job_store_redis_url() -> Optional[str]:
    return os.getenv("GRANTFLOW_JOB_STORE_REDIS_URL") or os.getenv("GRANTFLOW_REDIS_URL")


class RedisJobStore:
    """Redis-backed job store. Falls back to in-process dict if Redis is unavailable."""

    def __init__(self) -> None:
        self._fallback = InMemoryJobStore()
        self._redis: Any = None
        self._redis_lock = threading.Lock()
        self._url = _job_store_redis_url()

    def _client(self) -> Any:
        if self._redis is not None:
            return self._redis
        if not self._url:
            return None
        with self._redis_lock:
            if self._redis is not None:
                return self._redis
            try:
                import redis as _redis_lib

                client = _redis_lib.from_url(self._url, decode_responses=False, socket_connect_timeout=2)
                client.ping()
                self._redis = client
                return client
            except Exception:
                return None

    def _key(self, job_id: str) -> str:
        return f"{_REDIS_JOB_KEY_PREFIX}{job_id}"

    def set(self, job_id: str, payload: Dict[str, Any]) -> None:
        r = self._client()
        if r is None:
            self._fallback.set(job_id, payload)
            return
        stored = prepare_job_payload_for_storage(payload)
        try:
            r.setex(self._key(job_id), _REDIS_JOB_TTL_SECONDS, storage_json_dumps(stored).encode())
            r.sadd(_REDIS_JOB_INDEX_KEY, job_id)
            r.expire(_REDIS_JOB_INDEX_KEY, _REDIS_JOB_TTL_SECONDS)
        except Exception:
            self._fallback.set(job_id, payload)

    def update(self, job_id: str, **patch: Any) -> Dict[str, Any]:
        r = self._client()
        if r is None:
            return self._fallback.update(job_id, **patch)
        try:
            raw = r.get(self._key(job_id))
            current: Dict[str, Any] = storage_json_loads(raw.decode()) if raw else {}
            merged = dict(current)
            merged.update(prepare_job_payload_for_storage(patch))
            r.setex(self._key(job_id), _REDIS_JOB_TTL_SECONDS, storage_json_dumps(merged).encode())
            r.sadd(_REDIS_JOB_INDEX_KEY, job_id)
            r.expire(_REDIS_JOB_INDEX_KEY, _REDIS_JOB_TTL_SECONDS)
            return restore_job_payload_from_storage(merged)
        except Exception:
            return self._fallback.update(job_id, **patch)

    def get(self, job_id: str) -> Optional[Dict[str, Any]]:
        r = self._client()
        if r is None:
            return self._fallback.get(job_id)
        try:
            raw = r.get(self._key(job_id))
            if not raw:
                return None
            return restore_job_payload_from_storage(storage_json_loads(raw.decode()))
        except Exception:
            return self._fallback.get(job_id)

    def list(self) -> Dict[str, Dict[str, Any]]:
        r = self._client()
        if r is None:
            return self._fallback.list()
        try:
            job_ids = r.smembers(_REDIS_JOB_INDEX_KEY)
            result: Dict[str, Dict[str, Any]] = {}
            for jid_bytes in job_ids:
                jid = jid_bytes.decode() if isinstance(jid_bytes, bytes) else str(jid_bytes)
                raw = r.get(self._key(jid))
                if raw:
                    result[jid] = restore_job_payload_from_storage(storage_json_loads(raw.decode()))
            return result
        except Exception:
            return self._fallback.list()


def create_job_store_from_env() -> InMemoryJobStore | SQLiteJobStore | RedisJobStore:
    mode = storage_mode("GRANTFLOW_JOB_STORE", _env("JOB_STORE", "inmem"))
    if mode == "sqlite":
        return SQLiteJobStore()
    if mode == "redis":
        return RedisJobStore()
    return InMemoryJobStore()


def create_ingest_audit_store_from_env() -> InMemoryIngestAuditStore | SQLiteIngestAuditStore:
    mode = storage_mode(
        "GRANTFLOW_INGEST_STORE",
        _env("INGEST_STORE", storage_mode("GRANTFLOW_JOB_STORE", _env("JOB_STORE", "inmem"))),
    )
    if mode == "sqlite":
        return SQLiteIngestAuditStore()
    return InMemoryIngestAuditStore()
