#!/usr/bin/env python3
import io
import json
import logging
import os
import signal
import subprocess
import threading
import time
from datetime import datetime
from typing import Optional, Dict, Any, List, Union

import psutil
from pubsub import pub
from sqlalchemy import Column, Integer, String, DateTime, JSON
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.ext.mutable import MutableDict

from src.utilities.misc_helper import get_time_elapsed
from src.utilities.status_utils import RenderStatus, string_to_status

logger = logging.getLogger()
Base = declarative_base()


class BaseRenderWorker(Base):
    __tablename__ = 'render_workers'

    id = Column(String, primary_key=True)
    hostname = Column(String, nullable=True)
    input_path = Column(String)
    output_path = Column(String)
    date_created = Column(DateTime)
    start_time = Column(DateTime, nullable=True)
    end_time = Column(DateTime, nullable=True)
    engine_name = Column(String)
    engine_version = Column(String)
    engine_path = Column(String)
    priority = Column(Integer)
    start_frame = Column(Integer)
    end_frame = Column(Integer, nullable=True)
    parent = Column(String, nullable=True)
    children = Column(MutableDict.as_mutable(JSON))
    args = Column(MutableDict.as_mutable(JSON))
    name = Column(String)
    file_hash = Column(String)
    _status = Column(String)

    # --------------------------------------------
    # Required Overrides for Subclasses:
    # --------------------------------------------

    def __init__(self, input_path: str, output_path: str, engine_path: str, priority: int = 2,
                 args: Optional[Dict[str, Any]] = None, ignore_extensions: bool = True,
                 parent: Optional[str] = None, name: Optional[str] = None) -> None:
        """Initialize a render worker.

        Args:
            input_path: Path to the input project file.
            output_path: Path where output files will be saved.
            engine_path: Path to the render engine executable.
            priority: Job priority level (default: 2).
            args: Additional arguments for the render job.
            ignore_extensions: Whether to ignore file extension validation.
            parent: Parent job ID for distributed jobs.
            name: Custom name for the job.

        Raises:
            ValueError: If file extension is not supported.
            NotImplementedError: If engine is not defined.
        """
        if not ignore_extensions:
            if not any(ext in input_path for ext in self.engine.supported_extensions()):
                err_meg = f"Cannot find valid project with supported file extension for '{self.engine.name()}'"
                logger.error(err_meg)
                raise ValueError(err_meg)
        if not self.engine:
            raise NotImplementedError(f"Engine not defined for {self.__class__.__name__}")

        def generate_id():
            """Generate a unique job ID."""
            import uuid
            return str(uuid.uuid4()).split('-')[0]

        # Essential Info
        self.id = generate_id()
        self.hostname = None
        self.input_path = input_path
        self.output_path = output_path
        self.args = args or {}
        self.date_created = datetime.now()
        self.engine_name = self.engine.name()
        self.engine_path = engine_path
        self.engine_version = self.engine(engine_path).version()
        self.custom_engine_path = None
        self.priority = priority
        self.parent = parent
        self.children = {}
        self.name = name or os.path.basename(input_path)
        self.maximum_attempts = 3

        # Frame Ranges
        self.current_frame = 0
        self.start_frame = 0
        self.end_frame = None

        # Logging
        self.start_time = None
        self.end_time = None

        # History
        self.status = RenderStatus.CONFIGURING
        self.warnings = []
        self.errors = []

        # Threads and processes
        self.__thread = threading.Thread(target=self.__run, args=())
        self.__thread.daemon = True
        self.__process = None
        self.last_output = None
        self.__last_output_time = None
        self.watchdog_timeout = 120

    def generate_worker_subprocess(self) -> List[str]:
        """Generate a return a list of the command line arguments necessary to perform requested job

        Returns:
            list[str]: list of command line arguments
        """
        raise NotImplementedError("generate_worker_subprocess not implemented")

    def _parse_stdout(self, line: str) -> None:
        """Parses a line of standard output from the engine.

        This method should be overridden in a subclass to implement the logic for processing
        and interpreting a single line of output from the engine's standard output stream.

        On frame completion, the subclass should:
            1. Update value of self.current_frame
            2. Call self._send_frame_complete_notification()

        Args:
            line (str): A line of text from the engine's standard output.

        Raises:
            NotImplementedError: If the method is not overridden in a subclass.
        """
        raise NotImplementedError(f"_parse_stdout not implemented for {self.__class__.__name__}")

    # --------------------------------------------
    # Optional Overrides for Subclasses:
    # --------------------------------------------

    def percent_complete(self) -> float:
        """Return the completion percentage of this job.

        Returns:
            float: Completion percentage between 0.0 and 1.0.
        """
        # todo: fix this
        if self.status == RenderStatus.COMPLETED:
            return 1.0
        return 0

    def post_processing(self) -> None:
        """Override to perform any engine-specific postprocessing"""
        pass

    # --------------------------------------------
    # Do Not Override These Methods:
    # --------------------------------------------

    def __repr__(self):
        """Return string representation of the worker.

        Returns:
            str: String representation showing job details.
        """
        return f"<Job id:{self.id} p{self.priority} {self.engine_name}-{self.engine_version} '{self.name}' status:{self.status.value}>"

    @property
    def total_frames(self):
        return max(self.end_frame, 1) - self.start_frame + 1

    @property
    def status(self):
        return self._status

    @status.setter
    def status(self, new_status):
        if self._status != new_status.value:
            old_status = self.status
            self._status = new_status.value
            pub.sendMessage('status_change', job_id=self.id, old_status=old_status, new_status=new_status)

    @status.getter
    def status(self):
        if self._status in [RenderStatus.RUNNING.value, RenderStatus.NOT_STARTED.value]:
            if not hasattr(self, 'errors'):
                self._status = RenderStatus.CANCELLED.value
        return string_to_status(self._status)

    def _send_frame_complete_notification(self):
        pub.sendMessage('frame_complete', job_id=self.id, frame_number=self.current_frame)

    def generate_subprocess(self):
        """Generate subprocess command arguments.

        Returns:
            List[str]: List of command line arguments for the subprocess.

        Raises:
            ValueError: If argument conflicts are detected.
        """
        # Convert raw args from string if available and catch conflicts
        generated_args = [str(x) for x in self.generate_worker_subprocess()]
        generated_args_flags = [x for x in generated_args if x.startswith('-')]
        if len(generated_args_flags) != len(set(generated_args_flags)):
            msg = f"Cannot generate subprocess - Multiple arg conflicts detected: {generated_args}"
            logger.error(msg)
            raise ValueError(msg)
        return generated_args

    def get_raw_args(self):
        """Parse raw command line arguments from args dictionary.

        Returns:
            Optional[List[str]]: Parsed raw arguments, or None if no raw args.
        """
        raw_args_string = self.args.get('raw', '')
        raw_args = None
        if raw_args_string:
            import shlex
            raw_args = shlex.split(raw_args_string)
        return raw_args

    def log_path(self):
        """Generate the log file path for this job.

        Returns:
            str: Path to the log file.
        """
        filename = (self.name or os.path.basename(self.input_path)) + '_' + \
                   self.date_created.strftime("%Y.%m.%d_%H.%M.%S") + '.log'
        return os.path.join(os.path.dirname(self.input_path), filename)

    def start(self):
        """Start the render job.

        Validates input paths and engine availability, then starts the render thread.
        """
        if self.status not in [RenderStatus.SCHEDULED, RenderStatus.NOT_STARTED, RenderStatus.CONFIGURING]:
            logger.error(f"Trying to start job with status: {self.status}")
            return

        if not os.path.exists(self.input_path):
            self.status = RenderStatus.ERROR
            msg = f'Cannot find input path: {self.input_path}'
            logger.error(msg)
            self.errors.append(msg)
            return

        if not os.path.exists(self.engine_path):
            self.status = RenderStatus.ERROR
            msg = f'Cannot find render engine path for {self.engine.name()}'
            logger.error(msg)
            self.errors.append(msg)
            return

        self.status = RenderStatus.RUNNING if not self.children else RenderStatus.WAITING_FOR_SUBJOBS
        self.start_time = datetime.now()
        self.__thread.start()

    # handle multiple attempts at running subprocess
    def __run__subprocess_cycle(self, log_file):
        subprocess_cmds = self.generate_subprocess()
        initial_file_count = len(self.file_list())
        failed_attempts = 0

        log_file.write(f"Running command: {subprocess_cmds}\n")
        log_file.write('=' * 80 + '\n\n')

        while True:
            # Log attempt #
            if failed_attempts:
                if failed_attempts >= self.maximum_attempts:
                    err_msg = f"Maximum attempts exceeded ({self.maximum_attempts})"
                    logger.error(err_msg)
                    self.status = RenderStatus.ERROR
                    self.errors.append(err_msg)
                    return
                else:
                    log_file.write(f'\n{"=" * 20} Attempt #{failed_attempts + 1} {"=" * 20}\n\n')
                    logger.warning(f"Restarting render - Attempt #{failed_attempts + 1}")
                    self.status = RenderStatus.RUNNING

            return_code = self.__setup_and_run_process(log_file, subprocess_cmds)

            message = f"{'=' * 50}\n\n{self.engine.name()} render ended with code {return_code} " \
                      f"after {self.time_elapsed()}\n\n"
            log_file.write(message)

            # don't try again if we've been cancelled
            if self.status in [RenderStatus.CANCELLED, RenderStatus.ERROR]:
                return

            # if file output hasn't increased, return as error, otherwise restart process.
            file_count_has_increased = len(self.file_list()) > initial_file_count
            if (self.status == RenderStatus.RUNNING) and file_count_has_increased and not return_code:
                break

            if return_code:
                err_msg = f"{self.engine.name()} render failed with code {return_code}"
                logger.error(err_msg)
                self.errors.append(err_msg)

            # handle instances where engine exits ok but doesnt generate files
            if not return_code and not file_count_has_increased:
                err_msg = (f"{self.engine.name()} render exited ok, but file count has not increased. "
                           f"Count is still {len(self.file_list())}")
                log_file.write(f'Error: {err_msg}\n\n')
                self.errors.append(err_msg)

            # only count the attempt as failed if engine creates no output - reset counter on successful output
            failed_attempts = 0 if file_count_has_increased else failed_attempts + 1

    def __run__wait_for_subjobs(self, logfile):
        from src.distributed_job_manager import DistributedJobManager
        DistributedJobManager.wait_for_subjobs(parent_job=self)

    @staticmethod
    def log_and_print(message, log_file, level='info'):
        if level == 'debug':
            logger.debug(message)
        elif level == 'error':
            logger.error(message)
        else:
            logger.info(message)
        log_file.write(f"{message}\n")

    def __run(self):

        # Setup logging
        log_dir = os.path.dirname(self.log_path())
        os.makedirs(log_dir, exist_ok=True)

        with open(self.log_path(), "a") as log_file:

            self.log_and_print(f"{self.start_time.isoformat()} - Starting "
                               f"{self.engine.name()} {self.engine_version} render job for {self.name} "
                               f"({self.input_path})", log_file)
            log_file.write(f"\n")
            if not self.children:
                self.__run__subprocess_cycle(log_file)
            else:
                self.__run__wait_for_subjobs(log_file)

            # Validate Output - End if missing frames
            if self.status == RenderStatus.RUNNING:
                file_list_length = len(self.file_list())
                expected_list_length = (self.end_frame - self.start_frame + 1) if self.end_frame else 1

                msg = f"Frames: Expected ({expected_list_length}) vs actual ({file_list_length}) for {self}"
                self.log_and_print(msg, log_file, 'debug')

                if file_list_length not in (expected_list_length, 1):
                    msg = f"Missing frames: Expected ({expected_list_length}) vs actual ({file_list_length})"
                    self.log_and_print(msg, log_file, 'error')
                    self.errors.append(msg)
                    self.status = RenderStatus.ERROR
                    # todo: create new subjob to generate missing frames

            # cleanup and close if cancelled / error
            if self.status in [RenderStatus.CANCELLED, RenderStatus.ERROR]:
                self.end_time = datetime.now()
                message = f"{self.engine.name()} render ended with status '{self.status.value}' " \
                          f"after {self.time_elapsed()}"
                self.log_and_print(message, log_file)
                log_file.close()
                return

            # Post Render Work
            if not self.parent:
                logger.debug(f"Starting post-processing work for {self}")
                self.log_and_print(f"Starting post-processing work for {self}", log_file, 'debug')
                self.post_processing()
                self.log_and_print(f"Completed post-processing work for {self}", log_file, 'debug')

            self.status = RenderStatus.COMPLETED
            self.end_time = datetime.now()
            message = f"Render {self.name} completed successfully after {self.time_elapsed()}"
            self.log_and_print(message, log_file)

    def __setup_and_run_process(self, f, subprocess_cmds):

        def watchdog():
            logger.debug(f'Starting process watchdog for {self} with {self.watchdog_timeout}s timeout')
            while self.__process.poll() is None:
                time_since_last_update = time.time() - self.__last_output_time
                if time_since_last_update > self.watchdog_timeout:
                    logger.error(f"Process for {self} terminated due to exceeding timeout ({self.watchdog_timeout}s)")
                    self.__kill_process()
                    break
                # logger.debug(f'Watchdog for {self} - Time since last update: {time_since_last_update}')
                time.sleep(1)

            logger.debug(f'Stopping process watchdog for {self}')

        return_code = -1
        watchdog_thread = threading.Thread(target=watchdog)
        watchdog_thread.daemon = True

        try:
            # Start process and get updates
            if os.name == 'posix':  # linux / mac
                self.__process = subprocess.Popen(subprocess_cmds, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                                                  universal_newlines=False, preexec_fn=os.setsid)
            else:  # windows
                creationflags = subprocess.CREATE_NEW_PROCESS_GROUP | subprocess.CREATE_NO_WINDOW
                self.__process = subprocess.Popen(subprocess_cmds, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                                                  universal_newlines=False,
                                                  creationflags=creationflags)

            # Start watchdog
            self.__last_output_time = time.time()
            watchdog_thread.start()

            for c in io.TextIOWrapper(self.__process.stdout, encoding="utf-8"):  # or another encoding
                self.last_output = c.strip()
                self.__last_output_time = time.time()
                try:
                    f.write(c)
                    f.flush()
                    os.fsync(f.fileno())
                except Exception as e:
                    logger.error(f"Error saving log to disk: {e}")

                try:
                    self._parse_stdout(c.strip())
                except Exception as e:
                    logger.error(f'Error parsing stdout: {e}')

            f.write('\n')

            # Check return codes and process
            return_code = self.__process.wait()
        except Exception as e:
            message = f'Uncaught error running render process: {e}'
            f.write(message)
            logger.exception(message)
            self.__kill_process()

        # let watchdog end before continuing - prevents multiple watchdogs running when process restarts
        if watchdog_thread.is_alive():
            watchdog_thread.join()

        return return_code

    def __kill_process(self):
        try:
            if self.__process.poll():
                return
            logger.debug(f"Trying to kill process {self.__process}")
            self.__process.terminate()
            self.__process.kill()
            if os.name == 'posix':  # linux / macos
                os.killpg(os.getpgid(self.__process.pid), signal.SIGTERM)
                os.killpg(os.getpgid(self.__process.pid), signal.SIGKILL)
            else:  # windows
                parent = psutil.Process(self.__process.pid)
                for child in parent.children(recursive=True):
                    child.kill()
            self.__process.wait(timeout=5)
            logger.debug(f"Process ended with status {self.__process.poll()}")
        except (ProcessLookupError, AttributeError, psutil.NoSuchProcess):
            pass
        except Exception as e:
            logger.error(f"Error stopping the process: {e}")

    def is_running(self):
        """Check if the render job is currently running.

        Returns:
            bool: True if the job is running, False otherwise.
        """
        if hasattr(self, '__thread'):
            return self.__thread.is_alive()
        return False

    def log_error(self, error_line, halt_render=False):
        """Log an error and optionally halt the render.

        Args:
            error_line: Error message to log.
            halt_render: Whether to stop the render due to this error.
        """
        logger.error(error_line)
        self.errors.append(error_line)
        if halt_render:
            self.stop(is_error=True)

    def stop(self, is_error=False):
        """Stop the render job.

        Args:
            is_error: Whether this stop is due to an error.
        """
        logger.debug(f"Stopping {self}")

        # cleanup status
        if self.status in [RenderStatus.RUNNING, RenderStatus.NOT_STARTED, RenderStatus.SCHEDULED,
                           RenderStatus.CONFIGURING]:
            if is_error:
                err_message = self.errors[-1] if self.errors else 'Unknown error'
                logger.error(f"Halting render due to error: {err_message}")
                self.status = RenderStatus.ERROR
            else:
                self.status = RenderStatus.CANCELLED

        self.__kill_process()
        if self.is_running():  # allow the log files to close
            self.__thread.join(timeout=5)

    def time_elapsed(self):
        """Get elapsed time for this job.

        Returns:
            str: Formatted time elapsed string.
        """
        return get_time_elapsed(self.start_time, self.end_time)

    def file_list(self):
        """Get list of output files for this job.

        Returns:
            List[str]: List of output file paths.
        """
        try:
            job_dir = os.path.dirname(self.output_path)
            file_list = [
                os.path.join(job_dir, file)
                for file in os.listdir(job_dir)
                if not file.startswith('.')  # Ignore hidden files
            ]
            file_list.sort()
            return file_list
        except FileNotFoundError:
            return []

    def json(self):
        """Convert worker to JSON-serializable dictionary.

        Returns:
            Dict[str, Any]: Dictionary representation of worker data.
        """
        job_dict = {
            'id': self.id,
            'name': self.name,
            'hostname': self.hostname,
            'input_path': self.input_path,
            'output_path': self.output_path,
            'priority': self.priority,
            'parent': self.parent,
            'children': self.children,
            'date_created': self.date_created,
            'start_time': self.start_time,
            'end_time': self.end_time,
            'status': self.status.value,
            'file_hash': self.file_hash,
            'percent_complete': self.percent_complete(),
            'file_count': len(self.file_list()),
            'engine': self.engine_name,
            'engine_version': self.engine_version,
            'errors': getattr(self, 'errors', None),
            'start_frame': self.start_frame,
            'end_frame': self.end_frame,
            'total_frames': self.total_frames,
            'last_output': getattr(self, 'last_output', None),
            'log_path': self.log_path(),
            'args': self.args
        }

        # convert to json and back to auto-convert dates to iso format
        def date_serializer(o):
            """Serialize datetime objects to ISO format."""
            if isinstance(o, datetime):
                return o.isoformat()
            return None

        json_convert = json.dumps(job_dict, default=date_serializer)
        worker_json = json.loads(json_convert)
        return worker_json


def timecode_to_frames(timecode, frame_rate):
    """Convert timecode string to frame number.

    Args:
        timecode: Timecode in format HH:MM:SS:FF.
        frame_rate: Frames per second.

    Returns:
        int: Frame number corresponding to timecode.
    """
    e = [int(x) for x in timecode.split(':')]
    seconds = (((e[0] * 60) + e[1] * 60) + e[2])
    frames = (seconds * frame_rate) + e[-1] + 1
    return frames
