import json

import dataiku

from backend.models.chart import ChartPlan, ChartPlanPayload
from backend.utils.json_utils import extract_json_string
from backend.utils.llm_utils import add_history_to_completion
from backend.utils.logging_utils import get_logger

logger = get_logger(__name__)


class VisualizationService:
    @staticmethod
    def prepare_chart_plan(llm_id, messages, data, columns, trace=None):
        logger.info("Preparing chart plan")
        client = dataiku.api_client()
        project = client.get_default_project()
        comp = project.get_llm(llm_id).new_completion()
        chart_types_rules = {
            "none": "no chart is appropriate or supported.",
            "line": "the data shows trends or changes over an ordered variable (typically time on x-axis, numeric values on y-axis, possibly multiple series).",
            "pie": "the data represents a small number of categories where showing proportion or whole-part breakdown is best.",
            "bar": "the data compares categories or groups along a discrete x-axis (counts, averages, totals, etc.).",
            "heatmap": "the data is best represented by two categorical/ordinal axes with intensity or value at their intersections (e.g., correlations, frequency across two dimensions).",
        }
        chart_types_list = ", ".join(list(map(lambda ct: '"' + ct + '"', chart_types_rules.keys())))
        chart_rules = "".join(f"- {key} if {desc}\n" for key, desc in chart_types_rules.items())
        system_prompt = f"""
        You are a chart assistant. Focus on the last user query and the snippet of the data that the assistant retrieved to answer it 
        and any helpful earlier context (only if it clarifies intent of the last user query and its corresponding assistant answer).
        Then:
        - Choose exactly one of these chart types: {chart_types_list}.
        {chart_rules}
        - Do not invent other chart types.
        - Provide a short justification for your choice.
        - If you choose a chart (not "none"), also return the chart configuration using the **ChartPlan** Pydantic model.

        Output ONLY strict JSON (no prose, no markdown) with keys:
        - "chart_type": {chart_types_list}
        - "justification": string
        - "chart_data": null if chart_type == "none", otherwise an object matching the ChartPlan model.

        ChartPlan Python model definition:
        {json.dumps(ChartPlan.model_json_schema(), indent=2)}
        """.strip()
        comp.with_message(system_prompt, role="system")
        comp = add_history_to_completion(completion=comp, messages=messages)
        user_prompt = f"""
        columns: {columns}
        dataSample: {data}
        """.strip()
        comp.with_message(user_prompt, role="user")
        result = {
            "chart_type": "none",
            "justification": None,
            "chart_data": None,
        }
        logger.info(
            "Generating chart plan:\nLLM id=[%s]\nCompletion Query=%s\nCompletion Settings=%s\n",
            llm_id,
            json.dumps(comp.cq, indent=2, sort_keys=True),
            json.dumps(comp.settings, indent=2, sort_keys=True),
        )
        try:
            resp = comp.execute()
            if resp and resp.success:
                raw = resp.text.strip()
                logger.info("Chart Plan raw response:\nLLM id=[%s]\nCompletion Raw Response=%s", llm_id, raw)
                obj = json.loads(extract_json_string(raw))
                payload = ChartPlanPayload.model_validate(obj)
                logger.info(
                    f"Decided the best chart type is : {payload.chart_type} with justification: {payload.justification}"
                )
                if payload.chart_type == "none":
                    result = {
                        "chart_type": "none",
                        "justification": payload.justification,
                        "chart_data": None,
                    }
                else:
                    result = {
                        "chart_type": payload.chart_type,
                        "justification": payload.justification,
                        "chart_data": payload.chart_data.model_dump(),
                    }
                if resp.trace and trace is not None:
                    trace.append_trace(resp.trace)
                return result
            else:
                raise Exception(f"Failed to generate stories prompt: {resp._raw.get('errorMessage') if resp else ''}")
        except Exception as e:
            logger.exception(f"Failed to generate chart params {e}")
            return result
