# [DEF:TaskPersistenceModule:Module]
# @SEMANTICS: persistence, sqlite, sqlalchemy, task, storage
# @PURPOSE: Handles the persistence of tasks using SQLAlchemy and the tasks.db database.
# @LAYER: Core
# @RELATION: Used by TaskManager to save and load tasks.
# @INVARIANT: Database schema must match the TaskRecord model structure.

# [SECTION: IMPORTS]
from datetime import datetime
from typing import List, Optional, Dict, Any
import json

from sqlalchemy.orm import Session
from ...models.task import TaskRecord
from ..database import TasksSessionLocal
from .models import Task, TaskStatus, LogEntry
from ..logger import logger, belief_scope
# [/SECTION]

# [DEF:TaskPersistenceService:Class]
# @SEMANTICS: persistence, service, database, sqlalchemy
# @PURPOSE: Provides methods to save and load tasks from the tasks.db database using SQLAlchemy.
class TaskPersistenceService:
    # [DEF:__init__:Function]
    # @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
    # [/DEF:__init__:Function]

    # [DEF:persist_task:Function]
    # @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.
    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
                record.environment_id = task.params.get("environment_id") or task.params.get("source_env_id")
                record.started_at = task.started_at
                record.finished_at = task.finished_at
                record.params = task.params
                record.result = task.result
                
                # Store logs as JSON, converting datetime to string
                record.logs = []
                for log in task.logs:
                    log_dict = log.dict()
                    if isinstance(log_dict.get('timestamp'), datetime):
                        log_dict['timestamp'] = log_dict['timestamp'].isoformat()
                    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()
    # [/DEF:persist_task:Function]

    # [DEF:persist_tasks:Function]
    # @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.
    def persist_tasks(self, tasks: List[Task]) -> None:
        with belief_scope("TaskPersistenceService.persist_tasks"):
            for task in tasks:
                self.persist_task(task)
    # [/DEF:persist_tasks:Function]

    # [DEF:load_tasks:Function]
    # @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.
    def load_tasks(self, limit: int = 100, status: Optional[TaskStatus] = 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 = []
                        if record.logs:
                            for log_data in record.logs:
                                # Handle timestamp conversion if it's a string
                                if isinstance(log_data.get('timestamp'), str):
                                    log_data['timestamp'] = datetime.fromisoformat(log_data['timestamp'])
                                logs.append(LogEntry(**log_data))

                        task = Task(
                            id=record.id,
                            plugin_id=record.type,
                            status=TaskStatus(record.status),
                            started_at=record.started_at,
                            finished_at=record.finished_at,
                            params=record.params or {},
                            result=record.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()
    # [/DEF:load_tasks:Function]

    # [DEF:delete_tasks:Function]
    # @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.
    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()
    # [/DEF:delete_tasks:Function]

# [/DEF:TaskPersistenceService:Class]
# [/DEF:TaskPersistenceModule:Module]