import json
import logging
import os
import shutil
import uuid
import warnings
from collections.abc import Callable
from datetime import datetime
from pathlib import Path

from dhenara.agent.dsl.base import NodeInput
from dhenara.agent.dsl.base.context_registry import ExecutionContextRegistry
from dhenara.agent.dsl.events import EventBus, EventType
from dhenara.agent.observability.types import ObservabilitySettings
from dhenara.agent.run.registry import resource_config_registry
from dhenara.agent.types.data import AgentRunConfig, RunEnvParams
from dhenara.agent.utils.git import RunOutcomeRepository
from dhenara.agent.utils.io.artifact_manager import ArtifactManager
from dhenara.agent.utils.shared import get_project_identifier
from dhenara.ai.types.resource import ResourceConfig

logger = logging.getLogger(__name__)


class RunContext:
    """Manages a single execution run of an agent."""

    def __init__(
        self,
        # Primary config
        run_config: AgentRunConfig | None = None,
        # Legacy parameters (deprecated) - maintained for backward compatibility
        root_component_id: str | None = None,
        project_root: Path | None = None,
        run_root: Path | None = None,
        run_root_subpath: str | None = None,
        run_id: str | None = None,
        observability_settings: ObservabilitySettings | None = None,
        previous_run_id: str | None = None,
        start_hierarchy_path: str | None = None,
        # Non-config parameters
        initial_inputs: dict | None = None,
        input_source: Path | None = None,  # Static inputs
        # Execution ID
        # Do not confuse with this with run-id.  run-id is for the artifacts/ folder naming.
        # execution_id is a unique tracking ID for this execution
        # An autogenerated run_id will have timestamp info
        # Examples:
        # run_id : run_20250627_175140_9677c2
        # execution_id: dadex-c446215ea2ce4f8e
        execution_id: str | None = None,
    ):
        # Handle legacy parameters with deprecation warnings
        legacy_params_provided = any(
            [
                root_component_id is not None,
                project_root is not None,
                run_root is not None,
                run_root_subpath is not None,
                run_id is not None,
                observability_settings is not None,
                previous_run_id is not None,
                start_hierarchy_path is not None,
            ]
        )

        if legacy_params_provided and run_config is not None:
            warnings.warn(
                "Both run_config and individual parameters provided. Individual parameters will be ignored. "
                "Please use only run_config in future versions.",
                DeprecationWarning,
                stacklevel=2,
            )
        elif legacy_params_provided and run_config is None:
            warnings.warn(
                "Individual run configuration parameters are deprecated. "
                "Please use AgentRunConfig instead. "
                "Support for individual parameters will be removed in a future version.",
                DeprecationWarning,
                stacklevel=2,
            )

        # Build run_config from legacy parameters if not provided
        if run_config is None:
            if not observability_settings:
                # If observability_settings is not provided,
                # traces and metrics will be disabled and logs will be enabled
                observability_settings = ObservabilitySettings(
                    enable_tracing=False,
                    enable_metrics=False,
                    enable_logging=True,
                )

            run_config = AgentRunConfig(
                root_component_id=root_component_id,
                project_root=str(project_root),
                run_root=run_root,
                run_root_subpath=run_root_subpath,
                run_id=run_id,
                observability_settings=observability_settings,
                previous_run_id=previous_run_id,
                start_hierarchy_path=start_hierarchy_path,
                enable_outcome_repo=True,
            )

        self.run_config = run_config

        # Validate required parameters
        if not (self.run_config.root_component_id and self.run_config.project_root):
            raise ValueError(
                "root_component_id and project_root must be set either in run_config or as individual parameters"
            )

        # Set core properties from run_config
        self.root_component_id = self.run_config.root_component_id
        self.project_root = Path(str(self.run_config.project_root))
        self.observability_settings = self.run_config.observability_settings

        # Non-config properties
        self.execution_id = execution_id or f"dadex-{uuid.uuid4().hex[:16]}"
        self.input_source = input_source
        self.initial_inputs = initial_inputs or {}

        # Derived properties
        self.project_identifier = get_project_identifier(project_dir=self.project_root)
        self.run_root = Path(self.run_config.run_root) if self.run_config.run_root else self.project_root / "runs"
        self.run_root_subpath = self.run_config.run_root_subpath

        # Store re-run parameters from config
        self.run_id = self.run_config.run_id
        self.previous_run_id = self.run_config.previous_run_id
        self.start_hierarchy_path = self.run_config.start_hierarchy_path

        # Initialize core components
        self.execution_context_registry = ExecutionContextRegistry(enable_caching=True)
        self.event_bus = EventBus()
        self.setup_completed = False
        self.created_at = datetime.now()

        logger.info(f"Run context initialized with run configs: {self.run_config}")

    @property
    def effective_run_root(self) -> Path:
        """
        The actual root under which this run (and its outcomes) live:
        base run_root + optional subpath.
        """
        if self.run_root_subpath:
            return self.run_root / self.run_root_subpath
        return self.run_root

    def set_previous_run(
        self,
        previous_run_id: str,
        start_hierarchy_path: str | None = None,
    ):
        """Update previous run information. Updates both instance and run_config."""
        self.previous_run_id = previous_run_id
        self.start_hierarchy_path = start_hierarchy_path

        # Keep run_config in sync
        self.run_config.previous_run_id = previous_run_id
        self.run_config.start_hierarchy_path = start_hierarchy_path

    def setup_run(self, run_config: AgentRunConfig | None = None):
        """Setup the run environment. Uses instance run_config if none provided."""
        config = run_config or self.run_config

        # Indicates if this is a rerun of a previous execution
        self.is_rerun = self.previous_run_id is not None

        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        run_or_rerun = "rerun" if self.is_rerun else "run"
        _prefix = f"{config.run_id_prefix}_" if config.run_id_prefix else ""
        self.run_id = f"{_prefix}{run_or_rerun}_{timestamp}_{uuid.uuid4().hex[:6]}"
        self.run_dir = self.effective_run_root / self.run_id

        self.static_inputs_dir = self.run_dir / "static_inputs"
        self.static_inputs_dir.mkdir(parents=True, exist_ok=True)

        self.start_time = datetime.now()
        self.end_time = None
        self.metadata = {}

        # Create directories
        self.run_dir.mkdir(parents=True, exist_ok=True)
        # self.state_dir = self.run_dir / ".state"
        # self.state_dir.mkdir(parents=True, exist_ok=True)
        self.trace_dir = self.run_dir / ".trace"
        self.trace_dir.mkdir(parents=True, exist_ok=True)

        # Initialize git outcome repository
        if config.enable_outcome_repo:
            # Outcome is not inside the run id, there is a global outcome with
            # self.outcome_root = self.run_root
            self.outcome_dir = self.effective_run_root / "outcome"
            # Outcome is the final outcome git repo, not just node outputs
            _outcome_repo_name = self.project_identifier
            self.outcome_repo_dir = self.outcome_dir / _outcome_repo_name

            self.outcome_dir.mkdir(parents=True, exist_ok=True)
            self.outcome_repo_dir.mkdir(parents=True, exist_ok=True)

            self.outcome_repo = RunOutcomeRepository(self.outcome_repo_dir)
            self.git_branch_name = f"run/{self.run_id}"
            self.outcome_repo.create_run_branch(self.git_branch_name)
        else:
            self.outcome_dir = None
            self.outcome_repo_dir = None
            self.outcome_repo = None
            self.git_branch_name = None

        # Initialize previous run context
        self.previous_run_dir = None
        if self.previous_run_id:
            self.previous_run_dir = self.effective_run_root / self.previous_run_id
            if not self.previous_run_dir.exists():
                logger.error(f"Previous run directory does not exist: {self.previous_run_dir}")
                self.previous_run_id = None
                self.previous_run_dir = None

        # Setup observability with rerun info
        self.setup_observability()

        # Add rerun info to metadata
        if self.is_rerun:
            self.metadata["rerun_info"] = {
                "previous_run_id": self.previous_run_id,
                "start_hierarchy_path": self.start_hierarchy_path,
            }

        # Create run environment parameters
        self.run_env_params = RunEnvParams(
            run_id=self.run_id,
            run_dir=str(self.run_dir),
            run_root=str(self.run_root),
            run_root_subpath=str(self.run_root_subpath) if self.run_root_subpath else None,
            effective_run_root=str(self.effective_run_root),
            trace_dir=str(self.trace_dir),
            outcome_repo_dir=str(self.outcome_repo_dir) if self.outcome_repo_dir else None,
        )

        # Initialize artifact manager
        self.artifact_manager = ArtifactManager(
            run_env_params=self.run_env_params,
            outcome_repo=self.outcome_repo,
        )

        # Initialize resource config
        self.resource_config: ResourceConfig = self.get_resource_config()
        self.static_inputs = {}

        # Save initial metadata
        self._save_metadata()

        # Mark setup completed
        self.setup_completed = True

    def register_node_static_input(self, node_id: str, input_data: NodeInput):
        """Register static input for a node."""
        self.static_inputs[node_id] = input_data

    def register_event_handlers(self, handlers_map: dict[EventType, Callable]):
        """Register event handlers"""
        for event_type, handler in handlers_map.items():
            if not isinstance(event_type, EventType):
                raise ValueError(f"Illegal key {event_type}. Should be of type EventType")
            if not isinstance(handler, Callable):
                raise ValueError(f"Illegal handler for {event_type}. Should be of type Callable")

            self.event_bus.register(event_type, handler)

    def _save_metadata(self):
        """Save metadata about this run."""
        metadata = {
            "run_id": self.run_id,
            "created_at": self.created_at.isoformat(),
            "status": "initialized",
            **self.metadata,
        }
        with open(self.trace_dir / "dad_metadata.json", "w") as f:
            json.dump(metadata, f, indent=2)

    def copy_input_files(
        self,
        source: Path | None = None,
        files: list | None = None,
    ):
        """Prepare input data and files for the run."""
        input_source_path = source or self.input_source

        if not (input_source_path and input_source_path.exists()):
            logger.info(f"input_source_path {input_source_path} does not exist. No static input files copied")
            return

        # Save input data
        input_files = files or []

        if input_source_path:
            input_dir_path = input_source_path
            if input_dir_path.exists() and input_dir_path.is_dir():
                input_files += list(input_dir_path.glob("*"))

        # Copy input files if provided
        if input_files:
            for file_path in input_files:
                src = Path(file_path)
                if src.exists():
                    dst = self.static_inputs_dir / src.name
                    shutil.copy2(src, dst)

    def read_static_inputs(self):
        """Read initial inputs from the root input dir"""
        _input_file = self.static_inputs_dir / "static_inputs.json"

        if _input_file.exists():
            with open(_input_file) as f:
                try:
                    _data = json.load(f)

                    for node_id, static_input in _data.items():
                        self.register_node_static_input(node_id, static_input)

                    logger.info(f"Successfully loaded static inputs from {_input_file}")
                except Exception as e:
                    logger.exception(f"read_static_inputs: Error: {e}")
        else:
            logger.info(f"Static inputs are not initialized as no file exists in {_input_file}")

    async def complete_run(
        self,
        status="completed",
        error_msg: str | None = None,
    ):
        """Mark the run as complete and save final metadata."""
        if self.setup_completed:  # Failed after setup_run()
            self.end_time = datetime.now()
            self.metadata["status"] = status
            self.metadata["completed_at"] = self.end_time.isoformat()
            self.metadata["duration_seconds"] = (self.end_time - self.start_time).total_seconds()
            self._save_metadata()

            # Complete run in git repository
            if self.outcome_repo:
                self.outcome_repo.complete_run(
                    run_id=self.run_id,
                    status=status,
                    commit_outcome=True,
                )
        else:
            # TODO_FUTURE: Record to a global recording system?
            logger.error(f"Completed run even before run_context setup. Error: {error_msg}")

    async def load_from_previous_run(
        self,
        is_component: bool,
        hierarchy_path: str,
        copy_artifacts: bool = True,
    ) -> dict | None:
        """Copy artifacts from previous run up to the start node."""
        if not self.previous_run_dir or not self.previous_run_dir.exists():
            logger.error(f"Cannot copy artifacts: Previous run directory not found. Looked for {self.previous_run_dir}")
            return

        if copy_artifacts:
            await self._copy_previous_run_artifacts(
                hierarchy_path=hierarchy_path,
                is_component=is_component,
            )
            logger.info(f"Copied previous execution result artifacts for {hierarchy_path}")

        return self.load_previous_run_execution_result_dict(
            hierarchy_path=hierarchy_path,
            is_component=is_component,
        )

    async def _copy_previous_run_artifacts(
        self,
        hierarchy_path: str,
        is_component: bool,
    ):
        """Copy artifacts for a specific node/component from previous run."""
        if not self.previous_run_dir:
            return

        # Determine the hierarchy path for this node
        try:
            element_hier_dir = hierarchy_path.replace(".", os.sep)
        except Exception as e:
            logger.error(f"Error while deriving hierarchy for previous run artifacts: Error: {e}")
            return

        # Define source and target directories
        src_input_dir = self.previous_run_dir / element_hier_dir
        dst_input_dir = self.run_dir / element_hier_dir

        # Ensure target directory exists
        dst_input_dir.mkdir(parents=True, exist_ok=True)

        # NOTE: Currently there is nothing to copy for components
        if is_component:
            return

        if src_input_dir.exists():
            # Copy input, output, and outcome files
            for file_name in ["outcome.json", "result.json"]:
                src_file = src_input_dir / file_name
                if src_file.exists():
                    dst_file = dst_input_dir / file_name
                    try:
                        shutil.copy2(src_file, dst_file)
                        logger.debug(f"Copied {file_name} for {hierarchy_path}")
                    except Exception as e:
                        logger.warning(f"Failed to copy {file_name} for {hierarchy_path}: {e}")

    def load_previous_run_execution_result_dict(
        self,
        hierarchy_path: str,
        is_component: bool,
    ) -> dict | None:
        """Load execution results for a specific node/component from previous run."""
        if not self.previous_run_dir:
            return None

        try:
            element_hier_dir = hierarchy_path.replace(".", os.sep)
        except Exception as e:
            logger.error(f"Error while reloading results: element_hier_dir:{element_hier_dir}, Error: {e}")
            return None

        # NOTE: Currently there is nothing to load for components
        if is_component:
            return None

        # Define source directory
        src_input_dir = self.previous_run_dir / element_hier_dir
        result_file = src_input_dir / "result.json"

        if src_input_dir.exists() and result_file.exists():
            try:
                with open(result_file) as f:
                    _results = json.load(f)
                    return _results
            except Exception as e:
                logger.error(f"Failed to load results from {result_file}: {e}")
                return None

        return None

    def setup_observability(self):
        """Set up observability for the run context.

        Ensures a fresh logging setup per run so logs are written to this run's files.
        """
        from dhenara.agent.observability import configure_observability, reset_logging

        # Derive file paths inside the trace dir
        self.trace_file = self.trace_dir / "trace.jsonl"
        self.metrics_file = self.trace_dir / "metrics.jsonl"
        self.log_file = self.trace_dir / "logs.jsonl"

        # Ensure directory exists
        self.trace_dir.mkdir(parents=True, exist_ok=True)

        # Touch files to avoid permission issues inside exporter
        for file in [self.trace_file, self.log_file, self.metrics_file]:
            Path(file).touch(exist_ok=True)

        # Rerun metadata
        if self.is_rerun:
            # Modify the service name to indicate it's a rerun
            self.observability_settings.service_name = f"{self.observability_settings.service_name}-rerun"

            # Set additional tracing attributes for rerun info
            if self.previous_run_id:
                # These will be picked up by the tracing system
                os.environ["OTEL_RESOURCE_ATTRIBUTES"] = (
                    f"previous_run_id={self.previous_run_id},start_hierarchy_path={self.start_hierarchy_path},"
                )

        # Inject file paths
        self.observability_settings.trace_file_path = str(self.trace_file)
        self.observability_settings.metrics_file_path = str(self.metrics_file)
        self.observability_settings.log_file_path = str(self.log_file)

        # Force fresh logging (previous run may have initialized once)
        try:
            reset_logging()
        except Exception:
            logger.debug("reset_logging failed or unavailable; continuing")

        configure_observability(self.observability_settings)
        logger.info(f"Tracing enabled. Traces will be written to: {self.trace_file}")

    def get_resource_config(
        self,
        resource_profile="default",
    ):
        try:
            # Get resource configuration from registry
            resource_config = resource_config_registry.get(resource_profile)
            if not resource_config:
                # Fall back to creating a new one
                credentials_file = self.project_root / ".dhenara" / ".secrets" / ".credentials.yaml"
                if not credentials_file.exists():
                    credentials_file= os.getenv("DAD_DAI_CREDENTIALS_FILE", None)

                resource_config = self.load_default_resource_config(credentials_file)
                resource_config_registry.register(resource_profile, resource_config)

            return resource_config
        except Exception as e:
            raise ValueError(f"Error in resource setup: {e}")

    def load_default_resource_config(
        self,
        credentials_file,
    ):
        resource_config = ResourceConfig()
        resource_config.load_from_file(
            credentials_file=credentials_file,
            init_endpoints=True,
        )
        logger.debug(f"Loaded credentials from {credentials_file}")
        return resource_config

    def set_execution_context_caching(self, enabled: bool = True):
        """
        Enable or disable context caching to optimize memory usage.

        When disabled, only path relationships are maintained, but actual context
        objects are not stored, significantly reducing memory usage for complex flows.
        """
        self.execution_context_registry.set_caching_enabled(enabled)

    def get_dad_template_static_variables(self) -> dict:
        """Get static variables from run environment parameters."""
        # Guaranteed vars
        variables = {
            # --- Externally exposed vars
            #    1.environment variables
            "run_id": self.run_env_params.run_id,
            "run_dir": self.run_env_params.run_dir,
            "run_root": self.run_env_params.run_root,
            "effective_run_root": self.run_env_params.effective_run_root,
            # --- Internal vars
            #    1. state variables
            # "_dad_trace_dir": str(self.run_env_params.trace_dir),
        }

        # Optional vars
        if self.run_env_params.outcome_repo_dir:
            variables["outcome_repo_dir"] = str(self.run_env_params.outcome_repo_dir)

        return variables
