"""
Results Management System for EMDS Workflows

Handles storage, retrieval, and management of workflow execution results
using SQLite database with spatial geometry support.

Author: EMDS QGIS Port POC
Date: 2025-11-07
"""

import sqlite3
import json
from pathlib import Path
from dataclasses import dataclass
from typing import Dict, List, Optional, Any, Tuple
from datetime import datetime
import pandas as pd
from enum import Enum


class ExecutionStatus(Enum):
    """Workflow execution status"""
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


@dataclass
class ExecutionResult:
    """Single workflow execution result"""
    execution_id: str
    workflow_step_id: str
    status: ExecutionStatus
    start_time: str
    end_time: Optional[str] = None
    duration_ms: int = 0
    rows_processed: int = 0
    error_message: Optional[str] = None
    results_data: Optional[Dict[str, Any]] = None
    metadata: Optional[Dict[str, Any]] = None

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary"""
        return {
            'execution_id': self.execution_id,
            'workflow_step_id': self.workflow_step_id,
            'status': self.status.value,
            'start_time': self.start_time,
            'end_time': self.end_time,
            'duration_ms': self.duration_ms,
            'rows_processed': self.rows_processed,
            'error_message': self.error_message,
            'results_data': self.results_data,
            'metadata': self.metadata
        }


class ResultsManager:
    """Manages execution results in SQLite database"""

    def __init__(self, db_path: str = None):
        """
        Initialize results manager

        Args:
            db_path: Path to SQLite database (default: ~/.emds/results.db)
        """
        if db_path is None:
            db_path = str(Path.home() / '.emds' / 'results.db')

        self.db_path = Path(db_path)
        self.db_path.parent.mkdir(parents=True, exist_ok=True)

        # Initialize database
        self._init_database()

    def _init_database(self) -> None:
        """Initialize database schema"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()

        # Create executions table
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS executions (
                execution_id TEXT PRIMARY KEY,
                workflow_step_id TEXT NOT NULL,
                status TEXT NOT NULL,
                start_time TEXT NOT NULL,
                end_time TEXT,
                duration_ms INTEGER DEFAULT 0,
                rows_processed INTEGER DEFAULT 0,
                error_message TEXT,
                created_at TEXT DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (workflow_step_id) REFERENCES workflow_steps(id)
            )
        ''')

        # Create results table (stores result columns and values)
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS results (
                result_id INTEGER PRIMARY KEY AUTOINCREMENT,
                execution_id TEXT NOT NULL,
                row_id INTEGER NOT NULL,
                column_name TEXT NOT NULL,
                value TEXT,
                FOREIGN KEY (execution_id) REFERENCES executions(execution_id)
            )
        ''')

        # Create metadata table (stores JSON metadata)
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS metadata (
                metadata_id INTEGER PRIMARY KEY AUTOINCREMENT,
                execution_id TEXT NOT NULL,
                key TEXT NOT NULL,
                value TEXT,
                FOREIGN KEY (execution_id) REFERENCES executions(execution_id)
            )
        ''')

        # Create indices
        cursor.execute('CREATE INDEX IF NOT EXISTS idx_workflow_step ON executions(workflow_step_id)')
        cursor.execute('CREATE INDEX IF NOT EXISTS idx_execution_id ON results(execution_id)')

        conn.commit()
        conn.close()

    def save_execution(self, result: ExecutionResult) -> bool:
        """
        Save execution result to database

        Args:
            result: ExecutionResult object

        Returns:
            True if successful
        """
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()

            # Save execution record
            cursor.execute('''
                INSERT OR REPLACE INTO executions
                (execution_id, workflow_step_id, status, start_time, end_time,
                 duration_ms, rows_processed, error_message)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
            ''', (
                result.execution_id,
                result.workflow_step_id,
                result.status.value,
                result.start_time,
                result.end_time,
                result.duration_ms,
                result.rows_processed,
                result.error_message
            ))

            # Save metadata if provided
            if result.metadata:
                for key, value in result.metadata.items():
                    cursor.execute('''
                        INSERT INTO metadata (execution_id, key, value)
                        VALUES (?, ?, ?)
                    ''', (result.execution_id, key, json.dumps(value)))

            conn.commit()
            conn.close()
            return True

        except Exception as e:
            print(f"Error saving execution: {e}")
            return False

    def save_results_dataframe(self, execution_id: str, results_df: pd.DataFrame) -> bool:
        """
        Save results DataFrame to database

        Args:
            execution_id: Execution ID
            results_df: Results DataFrame

        Returns:
            True if successful
        """
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()

            # Batch insert all values using executemany for performance.
            # For 10,000 rows x 30 columns this reduces 300,000 individual
            # INSERT calls down to a single executemany with a list of tuples.
            rows_to_insert = []
            for row_id, row in results_df.iterrows():
                for col_name, value in row.items():
                    value_str = str(value) if value is not None else None
                    rows_to_insert.append((execution_id, row_id, col_name, value_str))

            cursor.executemany(
                'INSERT INTO results (execution_id, row_id, column_name, value) VALUES (?, ?, ?, ?)',
                rows_to_insert,
            )

            conn.commit()
            conn.close()
            return True

        except Exception as e:
            print(f"Error saving results dataframe: {e}")
            return False

    def get_execution(self, execution_id: str) -> Optional[ExecutionResult]:
        """
        Retrieve execution result by ID

        Args:
            execution_id: Execution ID

        Returns:
            ExecutionResult or None
        """
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()

            cursor.execute('''
                SELECT execution_id, workflow_step_id, status, start_time, end_time,
                       duration_ms, rows_processed, error_message
                FROM executions
                WHERE execution_id = ?
            ''', (execution_id,))

            row = cursor.fetchone()
            conn.close()

            if not row:
                return None

            return ExecutionResult(
                execution_id=row[0],
                workflow_step_id=row[1],
                status=ExecutionStatus(row[2]),
                start_time=row[3],
                end_time=row[4],
                duration_ms=row[5],
                rows_processed=row[6],
                error_message=row[7]
            )

        except Exception as e:
            print(f"Error retrieving execution: {e}")
            return None

    def get_execution_results_dataframe(self, execution_id: str) -> Optional[pd.DataFrame]:
        """
        Retrieve results DataFrame for an execution

        Args:
            execution_id: Execution ID

        Returns:
            Results DataFrame or None
        """
        try:
            conn = sqlite3.connect(self.db_path)

            query = '''
                SELECT row_id, column_name, value
                FROM results
                WHERE execution_id = ?
                ORDER BY row_id, column_name
            '''

            df = pd.read_sql_query(query, conn, params=(execution_id,))
            conn.close()

            if df.empty:
                return None

            # Pivot to get original DataFrame structure
            pivot_df = df.pivot_table(
                index='row_id',
                columns='column_name',
                values='value',
                aggfunc='first'
            )

            return pivot_df

        except Exception as e:
            print(f"Error retrieving results dataframe: {e}")
            return None

    def get_workflow_executions(self, workflow_step_id: str) -> List[ExecutionResult]:
        """
        Get all executions for a workflow step

        Args:
            workflow_step_id: Workflow step ID

        Returns:
            List of ExecutionResult objects
        """
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()

            cursor.execute('''
                SELECT execution_id, workflow_step_id, status, start_time, end_time,
                       duration_ms, rows_processed, error_message
                FROM executions
                WHERE workflow_step_id = ?
                ORDER BY start_time DESC
            ''', (workflow_step_id,))

            rows = cursor.fetchall()
            conn.close()

            results = []
            for row in rows:
                results.append(ExecutionResult(
                    execution_id=row[0],
                    workflow_step_id=row[1],
                    status=ExecutionStatus(row[2]),
                    start_time=row[3],
                    end_time=row[4],
                    duration_ms=row[5],
                    rows_processed=row[6],
                    error_message=row[7]
                ))

            return results

        except Exception as e:
            print(f"Error retrieving workflow executions: {e}")
            return []

    def delete_execution(self, execution_id: str) -> bool:
        """
        Delete execution and associated results

        Args:
            execution_id: Execution ID

        Returns:
            True if successful
        """
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()

            # Delete results
            cursor.execute('DELETE FROM results WHERE execution_id = ?', (execution_id,))

            # Delete metadata
            cursor.execute('DELETE FROM metadata WHERE execution_id = ?', (execution_id,))

            # Delete execution
            cursor.execute('DELETE FROM executions WHERE execution_id = ?', (execution_id,))

            conn.commit()
            conn.close()
            return True

        except Exception as e:
            print(f"Error deleting execution: {e}")
            return False

    def clear_old_executions(self, days: int = 30) -> int:
        """
        Delete executions older than specified days

        Args:
            days: Number of days to keep

        Returns:
            Number of executions deleted
        """
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()

            # Get old execution IDs
            cursor.execute('''
                SELECT execution_id FROM executions
                WHERE datetime(start_time) < datetime('now', ? || ' days')
            ''', (f'-{days}',))

            old_ids = [row[0] for row in cursor.fetchall()]

            # Delete them
            for exec_id in old_ids:
                self.delete_execution(exec_id)

            conn.close()
            return len(old_ids)

        except Exception as e:
            print(f"Error clearing old executions: {e}")
            return 0

    def get_execution_history(self, limit: int = 50) -> List[ExecutionResult]:
        """
        Get recent execution history

        Args:
            limit: Maximum number of results

        Returns:
            List of ExecutionResult objects
        """
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()

            cursor.execute('''
                SELECT execution_id, workflow_step_id, status, start_time, end_time,
                       duration_ms, rows_processed, error_message
                FROM executions
                ORDER BY start_time DESC
                LIMIT ?
            ''', (limit,))

            rows = cursor.fetchall()
            conn.close()

            results = []
            for row in rows:
                results.append(ExecutionResult(
                    execution_id=row[0],
                    workflow_step_id=row[1],
                    status=ExecutionStatus(row[2]),
                    start_time=row[3],
                    end_time=row[4],
                    duration_ms=row[5],
                    rows_processed=row[6],
                    error_message=row[7]
                ))

            return results

        except Exception as e:
            print(f"Error retrieving execution history: {e}")
            return []
