"""
Business logic for trace operations.
Orchestrates parsing, repository, and external data source access.
"""
import json
import time
from typing import Any, Dict, List, Optional

import dataiku

from traces_explorer.backend.models.trace import Trace, TokenUsage
from traces_explorer.backend.repositories.trace_repository import TraceRepository
from traces_explorer.backend.services.pasted_trace_storage import PastedTraceStorage
from traces_explorer.backend.services.trace_parser import TraceParser


class TraceService:
    """
    Service layer for trace operations.
    Handles business logic and orchestrates repository and parser.
    """

    def __init__(self, repository: TraceRepository, parser: TraceParser, dataiku_api, logger):
        """
        Initialize trace service.

        Args:
            repository: TraceRepository for data access
            parser: TraceParser for parsing logic
            dataiku_api: Dataiku API client
            logger: Logger instance
        """
        self.repository = repository
        self.parser = parser
        self.dataiku_api = dataiku_api
        self.logger = logger
        self.pasted_storage = PastedTraceStorage(
            logger=logger, dataiku_api=dataiku_api)

    def _read_column_content(self, dataset_name: str, column_name: str) -> List[str]:
        """
        Read column data from Dataiku dataset.

        Args:
            dataset_name: Name of dataset
            column_name: Name of column

        Returns:
            List of column values as strings

        Raises:
            Exception: If dataset or column not found
        """
        try:
            dataset = dataiku.Dataset(dataset_name)
            data_frame = dataset.get_dataframe()
            return data_frame[column_name].tolist()
        except KeyError as e:
            self.logger.error(
                "Column '%s' not found in dataset '%s': %s",
                column_name,
                dataset_name,
                e,
            )
            raise Exception(
                f"Column '{column_name}' not found in dataset '{dataset_name}'"
            ) from e
        except Exception as e:
            self.logger.error("Error reading dataset '%s': %s", dataset_name, e)
            raise

    def _load_dataset_traces(self) -> List[Trace]:
        """
        Load traces from configured Dataiku dataset.

        Returns:
            Dataset traces
        """
        config = self.dataiku_api.webapp_config
        llm_response_column = config.get("llm_responses_column")
        llm_response_dataset = config.get("llm_responses_dataset")

        if not llm_response_dataset or not llm_response_column:
            self.logger.info("No dataset or column configured, skipping trace loading")
            return []

        try:
            responses = self._read_column_content(
                llm_response_dataset, llm_response_column
            )
        except Exception as e:
            self.logger.error("Failed to read traces from dataset: %s", e)
            return []

        dataset_traces: List[Trace] = []

        for response_index, response_str in enumerate(responses):
            try:
                trace_id = f"DATASET_{response_index}"
                trace_data = self.parser.parse_trace(response_str)
                trace_detail = self._create_trace_detail(
                    trace_id=trace_id,
                    trace_data=trace_data,
                    source="dataset",
                )
                dataset_traces.append(trace_detail)

            except json.JSONDecodeError as exc:
                self.logger.warn(
                    "Skipping response at index %d due to JSON parsing error: %s",
                    response_index,
                    exc,
                )
            except Exception as exc:
                self.logger.error(
                    "Unexpected error processing response at index %d: %s",
                    response_index,
                    exc,
                )

        return dataset_traces

    def _load_pasted_traces(self) -> List[Trace]:
        """
        Load all persisted pasted traces from workload storage into repository.

        Returns:
            List of Trace objects
        """
        traces_data = self.pasted_storage.list_traces()
        pasted_traces: List[Trace] = []

        for trace in traces_data:
            if not isinstance(trace, dict):
                self.logger.warn("Skipping invalid pasted trace: not a dict")
                continue

            try:
                pasted_traces.append(Trace.from_dict(trace, source="pasted"))
            except Exception as exc:
                self.logger.warn(
                    "Skipping invalid pasted trace trace '%s': %s",
                    trace.get("id"),
                    exc,
                )

        return pasted_traces

    def load_traces(self) -> int:
        """
        Load traces from dataset and pasted traces from workload storage.
        Clears repository first.

        Returns:
            Number of dataset traces loaded
        """
        self.repository.clear()

        dataset_traces = self._load_dataset_traces()
        if dataset_traces:
            self.repository.add_all(dataset_traces)

        pasted_traces = self._load_pasted_traces()
        if pasted_traces:
            self.repository.add_all(pasted_traces)

        self.logger.info(
            "Loaded %d dataset traces and %d pasted traces",
            len(dataset_traces),
            len(pasted_traces),
        )
        return len(dataset_traces) + len(pasted_traces)

    def _create_trace_detail(
        self,
        trace_id: str,
        trace_data: Dict[str, Any],
        source: str = "dataset",
    ) -> Trace:
        """
        Create Trace from parsed trace data.

        Helper method to avoid duplication between load and process operations.

        Args:
            trace_id: Unique identifier for the trace
            trace_data: Parsed trace data dictionary
            source: Source category of the trace

        Returns:
            Trace object with processed trace tree
        """
        trace_name = self.parser.extract_display_name(trace_data)
        result = self.parser.extract_output(trace_data)

        tokens_dict = {
            "promptTokens": 0,
            "completionTokens": 0,
            "totalTokens": 0,
            "estimatedCost": 0.0,
        }
        self.parser.accumulate_usage_metadata(trace_data, tokens_dict)

        tokens = TokenUsage(
            prompt_tokens=tokens_dict["promptTokens"],
            completion_tokens=tokens_dict["completionTokens"],
            total_tokens=tokens_dict["totalTokens"],
            estimated_cost=tokens_dict["estimatedCost"],
        )

        processed_node = self.parser.process_trace_add_ids(trace_data)

        return Trace(
            id=trace_id,
            name=trace_name,
            result=result,
            begin=trace_data.get("begin"),
            duration=trace_data.get("duration"),
            tokens=tokens,
            source=source,
            parent_node=processed_node,
        )

    def list_traces(self) -> List[Trace]:
        """
        List all traces.

        Returns:
            List of Trace objects
        """
        return self.repository.list_all()

    def get_trace(self, trace_id: str) -> Optional[Trace]:
        """
        Get trace by ID (full details).

        Args:
            trace_id: Trace ID to retrieve

        Returns:
            Trace if found, None otherwise
        """
        return self.repository.get_by_id(trace_id)

    def process_pasted_trace(self, trace_json: str) -> Trace:
        """
        Process and persist a pasted trace JSON string.

        Args:
            trace_json: Raw JSON string containing trace data

        Returns:
            Trace object with processed trace tree

        Raises:
            ValueError: If JSON is invalid or no trace found
        """
        try:
            trace_data = self.parser.parse_trace(trace_json)
        except json.JSONDecodeError as e:
            raise ValueError(f"Invalid JSON: {str(e)}") from e

        if not trace_data:
            raise ValueError("No valid trace object found in JSON")

        trace = self._create_trace_detail(
            trace_id=f"LOCAL_PASTED_TRACE_{int(time.time() * 1000)}",
            trace_data=trace_data,
            source="pasted",
        )

        # In no_storage mode, the current saved trace ID is returned to be
        # removed from repository to keep only the latest pasted trace.
        replaced_trace_id = self.pasted_storage.save_trace(trace)
        if replaced_trace_id:
            self.repository.remove(replaced_trace_id)
        self.repository.add(trace)

        return trace

    def delete_pasted_trace(self, trace_id: str) -> bool:
        """
        Delete one pasted trace from file storage and in-memory repository.

        Args:
            trace_id: ID of pasted trace

        Returns:
            True when trace existed in storage or repository, False otherwise
        """
        trace = self.repository.get_by_id(trace_id)
        existed_in_repo = bool(trace and trace.source == "pasted")

        deleted_from_storage = self.pasted_storage.delete_trace(trace_id)

        if existed_in_repo:
            self.repository.remove(trace_id)

        return existed_in_repo or deleted_from_storage

    def delete_all_pasted_traces(self) -> int:
        """
        Delete all pasted traces from file storage and repository.

        Returns:
            Number of deleted files
        """
        deleted_count = self.pasted_storage.delete_all()
        self.repository.remove_by_source("pasted")
        return deleted_count
