# #region TaskPersistenceModule [C:5] [TYPE Module] [SEMANTICS sqlalchemy, task, search, task-persistence-service]
# @BRIEF Handles the persistence of tasks using SQLAlchemy and the tasks.db database.
# @LAYER Core
# @PRE  Tasks database must be initialized with TaskRecord and TaskLogRecord schemas.
# @POST  Provides reliable storage and retrieval for task metadata and logs.
# @SIDE_EFFECT  Performs database I/O on tasks.db.
# @DATA_CONTRACT  Input[Task, LogEntry] -> Model[TaskRecord, TaskLogRecord]
# @RELATION DEPENDS_ON -> [TaskManager]
# @RELATION DEPENDS_ON -> [TaskGraph]
# @RELATION DEPENDS_ON -> [TasksSessionLocal]
# @INVARIANT  Database schema must match the TaskRecord model structure.
from datetime import UTC, datetime
import json
import re

from sqlalchemy.orm import Session

from ...models.mapping import Environment
from ...models.task import TaskLogRecord, TaskRecord
from ..database import TasksSessionLocal
from ..logger import belief_scope, logger
from .models import LogEntry, LogFilter, LogStats, Task, TaskLog, TaskStatus


# #region TaskPersistenceService [C:5] [TYPE Class] [SEMANTICS persistence, service, database, sqlalchemy]
# @BRIEF Provides methods to save, load, and delete task records in tasks.db using SQLAlchemy models.
# @PRE  TasksSessionLocal must provide an active SQLAlchemy session, Task inputs must expose id/plugin_id/status/params/result/logs fields, and TaskRecord plus Environment schemas must be available.
# @POST  Persist operations leave matching TaskRecord rows committed or rolled back without leaking sessions, load operations return reconstructed Task objects from stored TaskRecord rows, and delete operations remove only the addressed task rows.
# @SIDE_EFFECT  Opens SQLAlchemy sessions, reads and writes task_records rows, resolves environment foreign keys against environments, commits or rolls back transactions, and emits error logs on persistence failures.
# @DATA_CONTRACT  Input[Task | List[Task] | List[str] | Query(limit:int,status:Optional[TaskStatus])] -> Model[TaskRecord, Environment] -> Output[None | List[Task]]
# @RELATION DEPENDS_ON -> [TasksSessionLocal]
# @RELATION DEPENDS_ON -> [TaskRecord]
# @RELATION DEPENDS_ON -> [Environment]
# @RELATION DEPENDS_ON -> [TaskManager]
# @RELATION DEPENDS_ON -> [TaskGraph]
# @INVARIANT  Persistence must handle potentially missing task fields natively.
#
# @TEST_CONTRACT  TaskPersistenceContract ->
# {
#   required_fields: {},
#   invariants: [
#       "persist_task creates or updates a record",
#       "load_tasks retrieves valid Task instances",
#       "delete_tasks correctly removes records from the database"
#   ]
# }
# @TEST_FIXTURE  valid_task_persistence -> {"task_id": "123", "status": "PENDING"}
# @TEST_EDGE  persist_invalid_task_type -> raises Exception
# @TEST_EDGE  load_corrupt_json_params -> handled gracefully
# @TEST_INVARIANT  accurate_round_trip -> verifies: [valid_task_persistence, load_corrupt_json_params]
class TaskPersistenceService:
    # #region _json_load_if_needed [TYPE Function] [C:1]
    # @PURPOSE: Safely load JSON strings from DB if necessary
    # @PRE  value is an arbitrary database value
    # @POST  Returns parsed JSON object, list, string, or primitive
    @staticmethod
    def _json_load_if_needed(value):
        # Hot-path utility — no entry/exit logging to reduce noise
        if value is None:
            return None
        if isinstance(value, (dict, list)):
            return value
        if isinstance(value, str):
            stripped = value.strip()
            if stripped == "" or stripped.lower() == "null":
                return None
            try:
                return json.loads(stripped)
            except json.JSONDecodeError:
                return value
        return value
    # #endregion _json_load_if_needed
    # #region _parse_datetime [TYPE Function] [C:1]
    # @PURPOSE: Safely parse a datetime string from the database
    # @PRE  value is an ISO string or datetime object
    # @POST  Returns datetime object or None
    @staticmethod
    def _parse_datetime(value):
        # Hot-path utility — no entry/exit logging to reduce startup noise
        if value is None or isinstance(value, datetime):
            return value
        if isinstance(value, str):
            try:
                return datetime.fromisoformat(value)
            except ValueError:
                return None
        return None
    # #endregion _parse_datetime
    # #region _resolve_environment_id [TYPE Function] [C:3]
    # @PURPOSE: Resolve environment id into existing environments.id value to satisfy FK constraints.
    # @PRE  Session is active
    # @POST  Returns existing environments.id or None when unresolved.
    # @DATA_CONTRACT  Input[env_id: Optional[str]] -> Output[Optional[str]]
    # @RELATION DEPENDS_ON -> [Environment]
    @staticmethod
    def _resolve_environment_id(
        session: Session, env_id: str | None
    ) -> str | None:
        with belief_scope("_resolve_environment_id"):
            raw_value = str(env_id or "").strip()
            if not raw_value:
                return None
            # 1) Direct match by primary key.
            by_id = (
                session.query(Environment).filter(Environment.id == raw_value).first()
            )
            if by_id:
                return str(by_id.id)
            # 2) Exact match by name.
            by_name = (
                session.query(Environment).filter(Environment.name == raw_value).first()
            )
            if by_name:
                return str(by_name.id)
            # 3) Slug-like match (e.g. "ss-dev" -> "SS DEV").
            def normalize_token(value: str) -> str:
                lowered = str(value or "").strip().lower()
                return re.sub(r"[^a-z0-9]+", "-", lowered).strip("-")
            target_token = normalize_token(raw_value)
            if not target_token:
                return None
            for env in session.query(Environment).all():
                if (
                    normalize_token(env.id) == target_token
                    or normalize_token(env.name) == target_token
                ):
                    return str(env.id)
            return None
    # #endregion _resolve_environment_id
    # #region __init__ [TYPE Function] [C:3]
    # @PURPOSE: Initializes the persistence service.
    # @PRE      None.
    # @POST     Service is ready.
    def __init__(self):
        with belief_scope("TaskPersistenceService.__init__"):
            # We use TasksSessionLocal from database.py
            pass
    # #endregion __init__
    # #region persist_task [TYPE Function] [C:3]
    # @PURPOSE: Persists or updates a single task in the database.
    # @PRE      isinstance(task, Task)
    # @POST     Task record created or updated in database.
    # @PARAM    task (Task) - The task object to persist.
    # @SIDE_EFFECT  Writes to task_records table in tasks.db
    # @DATA_CONTRACT  Input[Task] -> Model[TaskRecord]
    # @RELATION CALLS -> [_resolve_environment_id]
    def persist_task(self, task: Task) -> None:
        with belief_scope("TaskPersistenceService.persist_task", f"task_id={task.id}"):
            session: Session = TasksSessionLocal()
            try:
                record = (
                    session.query(TaskRecord).filter(TaskRecord.id == task.id).first()
                )
                if not record:
                    record = TaskRecord(id=task.id)
                    session.add(record)
                record.type = task.plugin_id
                record.status = task.status.value
                raw_env_id = task.params.get("environment_id") or task.params.get(
                    "source_env_id"
                )
                record.environment_id = self._resolve_environment_id(
                    session, raw_env_id
                )
                record.started_at = task.started_at
                record.finished_at = task.finished_at
                # Ensure params and result are JSON serializable
                def json_serializable(obj):
                    # Hot-path utility — no entry/exit logging
                    if isinstance(obj, dict):
                        return {k: json_serializable(v) for k, v in obj.items()}
                    elif isinstance(obj, list):
                        return [json_serializable(v) for v in obj]
                    elif isinstance(obj, datetime):
                        return obj.isoformat()
                    return obj
                record.params = json_serializable(task.params)
                record.result = json_serializable(task.result)
                # Store logs as JSON, converting datetime to string
                record.logs = []
                for log in task.logs:
                    log_dict = log.model_dump()
                    if isinstance(log_dict.get("timestamp"), datetime):
                        log_dict["timestamp"] = log_dict["timestamp"].isoformat()
                    # Also clean up any datetimes in context
                    if log_dict.get("context"):
                        log_dict["context"] = json_serializable(log_dict["context"])
                    record.logs.append(log_dict)
                # Extract error if failed
                if task.status == TaskStatus.FAILED:
                    for log in reversed(task.logs):
                        if log.level == "ERROR":
                            record.error = log.message
                            break
                session.commit()
            except Exception as e:
                session.rollback()
                logger.error(f"Failed to persist task {task.id}: {e}")
            finally:
                session.close()
    # #endregion persist_task
    # #region persist_tasks [TYPE Function] [C:3]
    # @PURPOSE: Persists multiple tasks.
    # @PRE      isinstance(tasks, list)
    # @POST     All tasks in list are persisted.
    # @PARAM    tasks (List[Task]) - The list of tasks to persist.
    # @RELATION CALLS -> [persist_task]
    def persist_tasks(self, tasks: list[Task]) -> None:
        with belief_scope("TaskPersistenceService.persist_tasks"):
            for task in tasks:
                self.persist_task(task)
    # #endregion persist_tasks
    # #region load_tasks [TYPE Function] [C:3]
    # @PURPOSE: Loads tasks from the database.
    # @PRE      limit is an integer.
    # @POST     Returns list of Task objects.
    # @PARAM    limit (int) - Max tasks to load.
    # @PARAM    status (Optional[TaskStatus]) - Filter by status.
    # @RETURN   List[Task] - The loaded tasks.
    # @DATA_CONTRACT  Model[TaskRecord] -> Output[List[Task]]
    # @RELATION CALLS -> [_json_load_if_needed]
    # @RELATION CALLS -> [_parse_datetime]
    def load_tasks(
        self, limit: int = 100, status: TaskStatus | None = None
    ) -> list[Task]:
        with belief_scope("TaskPersistenceService.load_tasks"):
            session: Session = TasksSessionLocal()
            try:
                query = session.query(TaskRecord)
                if status:
                    query = query.filter(TaskRecord.status == status.value)
                records = (
                    query.order_by(TaskRecord.created_at.desc()).limit(limit).all()
                )
                loaded_tasks = []
                for record in records:
                    try:
                        logs = []
                        logs_payload = self._json_load_if_needed(record.logs)
                        if isinstance(logs_payload, list):
                            for log_data in logs_payload:
                                if not isinstance(log_data, dict):
                                    continue
                                log_data = dict(log_data)
                                log_data["timestamp"] = (
                                    self._parse_datetime(log_data.get("timestamp"))
                                    or datetime.now(UTC)
                                )
                                logs.append(LogEntry(**log_data))
                        started_at = self._parse_datetime(record.started_at)
                        finished_at = self._parse_datetime(record.finished_at)
                        params = self._json_load_if_needed(record.params)
                        result = self._json_load_if_needed(record.result)
                        task = Task(
                            id=record.id,
                            plugin_id=record.type,
                            status=TaskStatus(record.status),
                            started_at=started_at,
                            finished_at=finished_at,
                            params=params if isinstance(params, dict) else {},
                            result=result,
                            logs=logs,
                        )
                        loaded_tasks.append(task)
                    except Exception as e:
                        logger.error(f"Failed to reconstruct task {record.id}: {e}")
                return loaded_tasks
            finally:
                session.close()
    # #endregion load_tasks
    # #region delete_tasks [TYPE Function] [C:3]
    # @PURPOSE: Deletes specific tasks from the database.
    # @PRE      task_ids is a list of strings.
    # @POST     Specified task records deleted from database.
    # @PARAM    task_ids (List[str]) - List of task IDs to delete.
    # @SIDE_EFFECT  Deletes rows from task_records table.
    # @RELATION DEPENDS_ON -> [TaskRecord]
    def delete_tasks(self, task_ids: list[str]) -> None:
        if not task_ids:
            return
        with belief_scope("TaskPersistenceService.delete_tasks"):
            session: Session = TasksSessionLocal()
            try:
                session.query(TaskRecord).filter(TaskRecord.id.in_(task_ids)).delete(
                    synchronize_session=False
                )
                session.commit()
            except Exception as e:
                session.rollback()
                logger.error(f"Failed to delete tasks: {e}")
            finally:
                session.close()
    # #endregion delete_tasks
# #endregion TaskPersistenceService
# #region TaskLogPersistenceService [C:5] [TYPE Class] [SEMANTICS persistence, service, database, log, sqlalchemy]
# @BRIEF Provides methods to store, query, summarize, and delete task log rows in the task_logs table.
# @PRE  TasksSessionLocal must provide an active SQLAlchemy session, task_id inputs must identify task log rows, LogEntry batches must expose timestamp/level/source/message/metadata fields, and LogFilter inputs must provide pagination and filter attributes used by queries.
# @POST  add_logs commits all provided log entries or rolls back on failure, query methods return TaskLog or LogStats views reconstructed from TaskLogRecord rows, and delete methods remove only log rows matching the supplied task identifiers.
# @SIDE_EFFECT  Opens SQLAlchemy sessions, inserts, reads, aggregates, and deletes task_logs rows, serializes log metadata to JSON, commits or rolls back transactions, and emits error logs on persistence failures.
# @DATA_CONTRACT  Input[task_id:str, logs:List[LogEntry], log_filter:LogFilter, task_ids:List[str]] -> Model[TaskLogRecord] -> Output[None | List[TaskLog] | LogStats | List[str]]
# @RELATION DEPENDS_ON -> [TaskLogRecord]
# @RELATION DEPENDS_ON -> [TasksSessionLocal]
# @RELATION DEPENDS_ON -> [TaskManager]
# @RELATION DEPENDS_ON -> [EventBus]
# @INVARIANT  Log entries are batch-inserted for performance.
#
# @TEST_CONTRACT  TaskLogPersistenceContract ->
# {
#   required_fields: {},
#   invariants: [
#       "add_logs efficiently saves logs to the database",
#       "get_logs retrieves properly filtered LogEntry objects"
#   ]
# }
# @TEST_FIXTURE  valid_log_batch -> {"task_id": "123", "logs": [{"level": "INFO", "message": "msg"}]}
# @TEST_EDGE  empty_log_list -> no-op behavior
# @TEST_EDGE  add_logs_db_error -> rollback and log error
# @TEST_INVARIANT  accurate_log_aggregation -> verifies: [valid_log_batch]
class TaskLogPersistenceService:
    """
    Service for persisting and querying task logs.
    Supports batch inserts, filtering, and statistics.
    """
    # #region __init__ [TYPE Function] [C:3]
    # @PURPOSE: Initializes the TaskLogPersistenceService
    # @PRE  config is provided or defaults are used
    # @POST  Service is ready for log persistence
    def __init__(self, config=None):
        pass
    # #endregion __init__
    # #region add_logs [TYPE Function] [C:3]
    # @PURPOSE: Batch insert log entries for a task.
    # @PRE  logs is a list of LogEntry objects.
    # @POST  All logs inserted into task_logs table.
    # @PARAM  task_id (str) - The task ID.
    # @PARAM  logs (List[LogEntry]) - Log entries to insert.
    # @SIDE_EFFECT  Writes to task_logs table.
    # @DATA_CONTRACT  Input[List[LogEntry]] -> Model[TaskLogRecord]
    # @RELATION DEPENDS_ON -> [TaskLogRecord]
    def add_logs(self, task_id: str, logs: list[LogEntry]) -> None:
        if not logs:
            return
        with belief_scope("TaskLogPersistenceService.add_logs", f"task_id={task_id}"):
            session: Session = TasksSessionLocal()
            try:
                for log in logs:
                    record = TaskLogRecord(
                        task_id=task_id,
                        timestamp=log.timestamp,
                        level=log.level,
                        source=log.source or "system",
                        message=log.message,
                        metadata_json=json.dumps(log.metadata)
                        if log.metadata
                        else None,
                    )
                    session.add(record)
                session.commit()
            except Exception as e:
                session.rollback()
                logger.error(f"Failed to add logs for task {task_id}: {e}")
            finally:
                session.close()
    # #endregion add_logs
    # #region get_logs [TYPE Function] [C:3]
    # @PURPOSE: Query logs for a task with filtering and pagination.
    # @PRE  task_id is a valid task ID.
    # @POST  Returns list of TaskLog objects matching filters.
    # @PARAM  task_id (str) - The task ID.
    # @PARAM  log_filter (LogFilter) - Filter parameters.
    # @RETURN  List[TaskLog] - Filtered log entries.
    # @DATA_CONTRACT  Model[TaskLogRecord] -> Output[List[TaskLog]]
    # @RELATION DEPENDS_ON -> [TaskLogRecord]
    # @RELATION DEPENDS_ON -> [LogFilter]
    # @RELATION DEPENDS_ON -> [TaskLog]
    def get_logs(self, task_id: str, log_filter: LogFilter) -> list[TaskLog]:
        with belief_scope("TaskLogPersistenceService.get_logs", f"task_id={task_id}"):
            session: Session = TasksSessionLocal()
            try:
                query = session.query(TaskLogRecord).filter(
                    TaskLogRecord.task_id == task_id
                )
                # Apply filters
                if log_filter.level:
                    query = query.filter(
                        TaskLogRecord.level == log_filter.level.upper()
                    )
                if log_filter.source:
                    query = query.filter(TaskLogRecord.source == log_filter.source)
                if log_filter.search:
                    search_pattern = f"%{log_filter.search}%"
                    query = query.filter(TaskLogRecord.message.ilike(search_pattern))
                # Order by timestamp ascending (oldest first)
                query = query.order_by(TaskLogRecord.timestamp.asc())
                # Apply pagination
                records = query.offset(log_filter.offset).limit(log_filter.limit).all()
                logs = []
                for record in records:
                    metadata = None
                    if record.metadata_json:
                        try:
                            metadata = json.loads(record.metadata_json)
                        except json.JSONDecodeError:
                            metadata = None
                    logs.append(
                        TaskLog(
                            id=record.id,
                            task_id=record.task_id,
                            timestamp=record.timestamp,
                            level=record.level,
                            source=record.source,
                            message=record.message,
                            metadata=metadata,
                        )
                    )
                return logs
            finally:
                session.close()
    # #endregion get_logs
    # #region get_log_stats [TYPE Function] [C:3]
    # @PURPOSE: Get statistics about logs for a task.
    # @PRE  task_id is a valid task ID.
    # @POST  Returns LogStats with counts by level and source.
    # @PARAM  task_id (str) - The task ID.
    # @RETURN  LogStats - Statistics about task logs.
    # @DATA_CONTRACT  Model[TaskLogRecord] -> Output[LogStats]
    # @RELATION DEPENDS_ON -> [TaskLogRecord]
    # @RELATION DEPENDS_ON -> [LogStats]
    def get_log_stats(self, task_id: str) -> LogStats:
        with belief_scope(
            "TaskLogPersistenceService.get_log_stats", f"task_id={task_id}"
        ):
            session: Session = TasksSessionLocal()
            try:
                # Get total count
                total_count = (
                    session.query(TaskLogRecord)
                    .filter(TaskLogRecord.task_id == task_id)
                    .count()
                )
                # Get counts by level
                from sqlalchemy import func
                level_counts = (
                    session.query(TaskLogRecord.level, func.count(TaskLogRecord.id))
                    .filter(TaskLogRecord.task_id == task_id)
                    .group_by(TaskLogRecord.level)
                    .all()
                )
                by_level = {level: count for level, count in level_counts}
                # Get counts by source
                source_counts = (
                    session.query(TaskLogRecord.source, func.count(TaskLogRecord.id))
                    .filter(TaskLogRecord.task_id == task_id)
                    .group_by(TaskLogRecord.source)
                    .all()
                )
                by_source = {source: count for source, count in source_counts}
                return LogStats(
                    total_count=total_count, by_level=by_level, by_source=by_source
                )
            finally:
                session.close()
    # #endregion get_log_stats
    # #region get_sources [TYPE Function] [C:3]
    # @PURPOSE: Get unique sources for a task's logs.
    # @PRE  task_id is a valid task ID.
    # @POST  Returns list of unique source strings.
    # @PARAM  task_id (str) - The task ID.
    # @RETURN  List[str] - Unique source names.
    # @DATA_CONTRACT  Model[TaskLogRecord] -> Output[List[str]]
    # @RELATION DEPENDS_ON -> [TaskLogRecord]
    def get_sources(self, task_id: str) -> list[str]:
        with belief_scope(
            "TaskLogPersistenceService.get_sources", f"task_id={task_id}"
        ):
            session: Session = TasksSessionLocal()
            try:
                from sqlalchemy import distinct
                sources = (
                    session.query(distinct(TaskLogRecord.source))
                    .filter(TaskLogRecord.task_id == task_id)
                    .all()
                )
                return [s[0] for s in sources]
            finally:
                session.close()
    # #endregion get_sources
    # #region delete_logs_for_task [TYPE Function] [C:3]
    # @PURPOSE: Delete all logs for a specific task.
    # @PRE  task_id is a valid task ID.
    # @POST  All logs for the task are deleted.
    # @PARAM  task_id (str) - The task ID.
    # @SIDE_EFFECT  Deletes from task_logs table.
    # @RELATION DEPENDS_ON -> [TaskLogRecord]
    def delete_logs_for_task(self, task_id: str) -> None:
        with belief_scope(
            "TaskLogPersistenceService.delete_logs_for_task", f"task_id={task_id}"
        ):
            session: Session = TasksSessionLocal()
            try:
                session.query(TaskLogRecord).filter(
                    TaskLogRecord.task_id == task_id
                ).delete(synchronize_session=False)
                session.commit()
            except Exception as e:
                session.rollback()
                logger.error(f"Failed to delete logs for task {task_id}: {e}")
            finally:
                session.close()
    # #endregion delete_logs_for_task
    # #region delete_logs_for_tasks [TYPE Function] [C:3]
    # @PURPOSE: Delete all logs for multiple tasks.
    # @PRE  task_ids is a list of task IDs.
    # @POST  All logs for the tasks are deleted.
    # @PARAM  task_ids (List[str]) - List of task IDs.
    # @SIDE_EFFECT  Deletes rows from task_logs table.
    # @RELATION DEPENDS_ON -> [TaskLogRecord]
    def delete_logs_for_tasks(self, task_ids: list[str]) -> None:
        if not task_ids:
            return
        with belief_scope("TaskLogPersistenceService.delete_logs_for_tasks"):
            session: Session = TasksSessionLocal()
            try:
                session.query(TaskLogRecord).filter(
                    TaskLogRecord.task_id.in_(task_ids)
                ).delete(synchronize_session=False)
                session.commit()
            except Exception as e:
                session.rollback()
                logger.error(f"Failed to delete logs for tasks: {e}")
            finally:
                session.close()
    # #endregion delete_logs_for_tasks
# #endregion TaskLogPersistenceService
# #endregion TaskPersistenceModule
