# Copyright (c) 2025 CoReason, Inc.
#
# This software is proprietary and dual-licensed.
# Licensed under the Prosperity Public License 3.0 (the "License").
# A copy of the license is available at https://prosperitylicense.com/versions/3.0.0
# For details, see the LICENSE file.
# Commercial use beyond a 30-day trial requires a separate license.
#
# Source Code: https://github.com/CoReason-AI/coreason_prism

import io
from pathlib import Path
from typing import Any, Dict, List, Optional, Union, cast
from uuid import uuid4

import pandas as pd
from loguru import logger
from PIL import Image
from transformers import Pix2StructForConditionalGeneration, Pix2StructProcessor

from coreason_identity.models import UserContext
from coreason_prism.models.vision import ExtractedFigure, FigureType
from coreason_prism.utils.engine import HuggingFaceEngine


class Analyst(HuggingFaceEngine):
    """The Analyst (Chart Extraction Engine).

    Uses Vision Transformers (DePlot) to extract data from chart images,
    converting images of line graphs back into CSV data.
    """

    MODEL_NAME = "google/deplot"
    MODEL_CLASS = Pix2StructForConditionalGeneration
    PROCESSOR_CLASS = Pix2StructProcessor

    async def run(self, instruction: str, context: Dict[str, Any], *, user_context: UserContext) -> Dict[str, Any]:
        """Execute the agent task.

        Args:
            instruction (str): The task instruction.
            context (dict): The context containing data.
            user_context (UserContext): The user identity context.

        Returns:
            dict: The execution result.
        """
        logger.info(f"Analyst received instruction: {instruction}")
        image_path = context.get("image_path")
        if not image_path:
            raise ValueError("Analyst context missing 'image_path'")

        result = self.process_image(image_path, source_document_id="api_request", context=user_context)

        return {
            "status": "success",
            "content": f"Analyzed chart image. Type: {result.figure_type}",
            "artifacts": [{"type": "chart_data", "data": result.model_dump(mode="json")}],
        }

    @classmethod
    def process_image(
        cls, image_path: Union[str, Path], source_document_id: str, *, context: UserContext
    ) -> ExtractedFigure:
        """Process an image to extract underlying data using DePlot.

        Args:
            image_path (Union[str, Path]): Path to the image file.
            source_document_id (str): ID of the source document.
            context (UserContext): The user identity context.

        Returns:
            ExtractedFigure: A model containing the extracted data series and
            metadata (e.g., median survival time).

        Raises:
            FileNotFoundError: If the image path does not exist.
            Exception: For any model inference or processing errors.
        """
        path = Path(image_path)
        if not path.exists():
            logger.error(f"Image file not found: {path}")
            raise FileNotFoundError(f"Image file not found: {path}")

        cls._load_model()

        assert cls._model is not None
        assert cls._processor is not None
        assert cls._device is not None

        try:
            logger.info(
                "Starting DataFrame analysis",
                user_id=context.user_id.get_secret_value(),
            )
            logger.info(f"Processing image for data extraction: {path}")
            image = Image.open(path).convert("RGB")

            # DePlot prompt
            inputs = cls._processor(
                images=image, text="Generate underlying data table of the figure below:", return_tensors="pt"
            ).to(cls._device)

            predictions = cls._model.generate(**inputs, max_new_tokens=512)
            decoded_text = cast(Any, cls._processor).decode(predictions[0], skip_special_tokens=True)

            logger.debug(f"DePlot raw output: {decoded_text}")

            # 1. Parse to DataFrame
            df = cls._parse_deplot_output(decoded_text)

            # 2. Determine Type (using DF)
            figure_type = cls._determine_figure_type(df)

            # 3. Calculate Metadata (using DF)
            metadata: Dict[str, Any] = {}
            if not df.empty and df.shape[1] >= 2:
                median_survival = cls._calculate_median_survival(df)
                if median_survival is not None:
                    metadata["median_survival"] = median_survival
                    logger.info(f"Calculated median survival: {median_survival}")

            # 4. Convert to Dict for ExtractedFigure (Handling NaN -> None)
            # Replace NaNs with None for JSON compatibility
            # Ensure object dtype so None can be inserted
            if not df.empty:
                df_out = df.astype(object).where(pd.notnull(df), None)
                data_series = cast(Dict[str, List[Optional[float]]], df_out.to_dict(orient="list"))
            else:
                data_series = {}

            return ExtractedFigure(
                id=uuid4(),
                source_document_id=source_document_id,
                image_path=str(path),
                figure_type=figure_type,
                data_series=data_series,
                embedding=None,
                metadata=metadata if metadata else None,
            )

        except Exception as e:
            logger.exception(f"Error processing image {path}")
            raise e

    @staticmethod
    def _parse_deplot_output(text: str) -> pd.DataFrame:
        """Parse the linearized table output from DePlot into a pandas DataFrame.

        Args:
            text (str): The raw output string from the DePlot model.

        Returns:
            pd.DataFrame: A DataFrame containing the parsed data, or an empty
            DataFrame if parsing fails.
        """
        clean_text = text.replace("<0x0A>", "\n").replace("<s>", "").replace("</s>", "")

        if not clean_text.strip():
            logger.warning("DePlot produced empty output")
            return pd.DataFrame()

        try:
            df = pd.read_csv(
                io.StringIO(clean_text), sep="|", skipinitialspace=True, on_bad_lines="skip", index_col=False
            )

            # Strip whitespace from column names
            df.columns = df.columns.str.strip()

            # Filter out separator rows
            if not df.empty:
                first_row = df.iloc[0].astype(str)
                if first_row.str.match(r"^[-:\s]+$").all():
                    df = df.iloc[1:]

            # Convert to numeric, coercing errors to NaN
            for col in df.columns:
                df[col] = df[col].astype(str).str.replace(",", "").str.strip().str.strip("'\"")
                df[col] = pd.to_numeric(df[col], errors="coerce")

            return df

        except Exception as e:
            logger.warning(f"Pandas parsing failed, falling back to empty. Error: {e}")
            return pd.DataFrame()

    @staticmethod
    def _determine_figure_type(df: pd.DataFrame) -> FigureType:
        """Determine the figure type based on the extracted DataFrame structure.

        Args:
            df (pd.DataFrame): The extracted DataFrame.

        Returns:
            FigureType: The detected type (LINE_CHART, BAR_CHART, etc.).
        """
        if df.empty:
            return FigureType.UNKNOWN

        # Assume first column is X-axis
        x_col = df.iloc[:, 0]
        total_points = len(x_col)

        # Count numeric values (non-NaN)
        numeric_count = x_col.notna().sum()
        numeric_ratio = numeric_count / total_points

        if numeric_ratio > 0.8:
            return FigureType.LINE_CHART
        else:
            return FigureType.BAR_CHART

    @staticmethod
    def _calculate_median_survival(df: pd.DataFrame) -> Optional[float]:
        """Calculate the median survival (X where Y=50%) from the DataFrame.

        Used for Kaplan-Meier Survival Curves extraction.

        Args:
            df (pd.DataFrame): The DataFrame containing Time (X) and Survival % (Y).

        Returns:
            Optional[float]: The interpolated X value where Y is 50%, or None if
            it cannot be determined.
        """
        if df.shape[1] < 2:
            return None

        # Heuristic: Assume Col 0 is Time (X), Col 1 is Survival (Y)
        # Use .iloc for positional access
        x_col = df.columns[0]
        y_col = df.columns[1]

        # Filter out rows where X or Y is NaN
        # subset=[x_col, y_col] ensures we only care about these two
        clean_df = df.dropna(subset=[x_col, y_col]).copy()

        if len(clean_df) < 2:
            return None

        # Sort by Time (X)
        clean_df = clean_df.sort_values(by=x_col)

        # Get arrays for faster access
        xs = clean_df[x_col].values
        ys = clean_df[y_col].values

        # Determine scale
        max_y = ys.max()
        target_y = 50.0 if max_y > 1.0 else 0.5
        min_y = ys.min()

        if min_y > target_y:
            return None

        # Iterate through points to find crossing
        # We can use zip on the arrays
        for i in range(len(xs) - 1):
            x1, y1 = xs[i], ys[i]
            x2, y2 = xs[i + 1], ys[i + 1]

            if (y1 >= target_y >= y2) or (y1 <= target_y <= y2):
                if y1 == y2:
                    continue

                # Linear Interpolation
                fraction = (target_y - y1) / (y2 - y1)
                median_x = x1 + fraction * (x2 - x1)
                return float(median_x)

        return None
