import csv
from pathlib import Path
from typing import Any

from pydantic import Field, RootModel
from rich import box
from rich.table import Table

from pipelex import log
from pipelex.cogt.exceptions import CostRegistryError
from pipelex.cogt.llm.llm_report import LLMTokenCostReport, LLMTokenCostReportField, LLMTokensUsage
from pipelex.cogt.usage.cost_category import CostCategory, CostsByCategoryDict
from pipelex.cogt.usage.costs_per_token import model_cost_per_token
from pipelex.cogt.usage.token_category import TokenCategory
from pipelex.hub import get_console
from pipelex.tools.typing.pydantic_utils import empty_list_factory_of

CostRegistryRoot = list[LLMTokenCostReport]


class CostRegistry(RootModel[CostRegistryRoot]):
    root: CostRegistryRoot = Field(default_factory=empty_list_factory_of(LLMTokenCostReport))

    def to_records(self) -> list[dict[str, Any]]:
        """Convert cost reports to list of flat dictionaries."""
        records: list[dict[str, Any]] = []
        for token_cost_report in self.root:
            record_dict = token_cost_report.as_flat_dictionary()
            records.append(record_dict)
        return records

    @classmethod
    def generate_report(
        cls,
        pipeline_run_id: str,
        llm_tokens_usages: list[LLMTokensUsage],
        unit_scale: float,
        cost_report_file_path: Path | None = None,
    ):
        if not llm_tokens_usages:
            if pipeline_run_id != "untitled":
                log.warning(f"No report to generate for pipeline '{pipeline_run_id}'")
            else:
                log.verbose(f"No report to generate for pipeline '{pipeline_run_id}'")
            return
        cost_registry = CostRegistry()
        for llm_tokens_usage in llm_tokens_usages:
            cost_report = cls.complete_cost_report(llm_tokens_usage=llm_tokens_usage)
            cost_registry.root.append(cost_report)

        records = cost_registry.to_records()

        # Calculate total costs overall
        total_nb_tokens_input_cached = sum(record[LLMTokenCostReportField.NB_TOKENS_INPUT_CACHED] for record in records)
        total_nb_tokens_input_non_cached = sum(record[LLMTokenCostReportField.NB_TOKENS_INPUT_NON_CACHED] for record in records)
        total_nb_tokens_input_joined = sum(record[LLMTokenCostReportField.NB_TOKENS_INPUT_JOINED] for record in records)
        total_nb_tokens_output = sum(record[LLMTokenCostReportField.NB_TOKENS_OUTPUT] for record in records)
        total_cost_input_cached = sum(record[LLMTokenCostReportField.COST_INPUT_CACHED] for record in records)
        total_cost_input_non_cached = sum(record[LLMTokenCostReportField.COST_INPUT_NON_CACHED] for record in records)
        total_cost_input_joined = sum(record[LLMTokenCostReportField.COST_INPUT_JOINED] for record in records)
        total_cost_output = sum(record[LLMTokenCostReportField.COST_OUTPUT] for record in records)
        total_cost = cls.compute_total_cost(
            input_non_cached_cost=total_cost_input_non_cached,
            input_cached_cost=total_cost_input_cached,
            output_cost=total_cost_output,
        )

        # Calculate costs per LLM model - group by LLM name
        grouped_by_llm: dict[str, dict[str, float]] = {}
        for record in records:
            llm_name = record[LLMTokenCostReportField.LLM_NAME]
            if llm_name not in grouped_by_llm:
                grouped_by_llm[llm_name] = {
                    LLMTokenCostReportField.NB_TOKENS_INPUT_CACHED: 0,
                    LLMTokenCostReportField.NB_TOKENS_INPUT_NON_CACHED: 0,
                    LLMTokenCostReportField.NB_TOKENS_INPUT_JOINED: 0,
                    LLMTokenCostReportField.NB_TOKENS_OUTPUT: 0,
                    LLMTokenCostReportField.COST_INPUT_CACHED: 0.0,
                    LLMTokenCostReportField.COST_INPUT_NON_CACHED: 0.0,
                    LLMTokenCostReportField.COST_INPUT_JOINED: 0.0,
                    LLMTokenCostReportField.COST_OUTPUT: 0.0,
                }

            # Aggregate values
            for field in [
                LLMTokenCostReportField.NB_TOKENS_INPUT_CACHED,
                LLMTokenCostReportField.NB_TOKENS_INPUT_NON_CACHED,
                LLMTokenCostReportField.NB_TOKENS_INPUT_JOINED,
                LLMTokenCostReportField.NB_TOKENS_OUTPUT,
                LLMTokenCostReportField.COST_INPUT_CACHED,
                LLMTokenCostReportField.COST_INPUT_NON_CACHED,
                LLMTokenCostReportField.COST_INPUT_JOINED,
                LLMTokenCostReportField.COST_OUTPUT,
            ]:
                grouped_by_llm[llm_name][field] += record[field]

        if not grouped_by_llm:
            msg = "Empty report aggregation by LLM name"
            raise CostRegistryError(msg)

        console = get_console()
        title = f"Costs by LLM model for pipeline '{pipeline_run_id}'"
        table = Table(title=title, box=box.ROUNDED)

        scale_str: str
        if unit_scale == 1:
            scale_str = ""
        else:
            scale_str = str(unit_scale)
        # Add columns
        table.add_column("Model", style="cyan", overflow="fold", width=30)
        table.add_column("Input Cached", justify="right", style="green")
        table.add_column("Input Non Cached", justify="right", style="green")
        table.add_column("Input Joined", justify="right", style="green")
        table.add_column("Output", justify="right", style="green")
        table.add_column(f"Input Cached Cost ({scale_str}$)", justify="right", style="yellow")
        table.add_column(f"Input Non Cached Cost ({scale_str}$)", justify="right", style="yellow")
        table.add_column(f"Input Joined Cost ({scale_str}$)", justify="right", style="yellow")
        table.add_column(f"Output Cost ({scale_str}$)", justify="right", style="yellow")
        table.add_column(f"Total Cost ({scale_str}$)", justify="right", style="bold yellow")

        # Add rows for each LLM model
        for llm_name, aggregated_data in grouped_by_llm.items():
            row_total_cost = cls.compute_total_cost(
                input_non_cached_cost=aggregated_data[LLMTokenCostReportField.COST_INPUT_NON_CACHED],
                input_cached_cost=aggregated_data[LLMTokenCostReportField.COST_INPUT_CACHED],
                output_cost=aggregated_data[LLMTokenCostReportField.COST_OUTPUT],
            )
            table.add_row(
                llm_name,
                f"{int(aggregated_data[LLMTokenCostReportField.NB_TOKENS_INPUT_CACHED]):,}",
                f"{int(aggregated_data[LLMTokenCostReportField.NB_TOKENS_INPUT_NON_CACHED]):,}",
                f"{int(aggregated_data[LLMTokenCostReportField.NB_TOKENS_INPUT_JOINED]):,}",
                f"{int(aggregated_data[LLMTokenCostReportField.NB_TOKENS_OUTPUT]):,}",
                f"{aggregated_data[LLMTokenCostReportField.COST_INPUT_CACHED] / unit_scale:.4f}",
                f"{aggregated_data[LLMTokenCostReportField.COST_INPUT_NON_CACHED] / unit_scale:.4f}",
                f"{aggregated_data[LLMTokenCostReportField.COST_INPUT_JOINED] / unit_scale:.4f}",
                f"{aggregated_data[LLMTokenCostReportField.COST_OUTPUT] / unit_scale:.4f}",
                f"{row_total_cost / unit_scale:.4f}",
            )

        # add total row
        footer_style = "bold"
        table.add_row(
            "Total",
            f"{total_nb_tokens_input_cached:,}",
            f"{total_nb_tokens_input_non_cached:,}",
            f"{total_nb_tokens_input_joined:,}",
            f"{total_nb_tokens_output:,}",
            f"{total_cost_input_cached / unit_scale:.4f}",
            f"{total_cost_input_non_cached / unit_scale:.4f}",
            f"{total_cost_input_joined / unit_scale:.4f}",
            f"{total_cost_output / unit_scale:.4f}",
            f"{total_cost / unit_scale:.4f}",
            style=footer_style,
            end_section=True,
        )

        console.print(table)
        console.print(" [dim]Note: the report only includes LLMs and some costs might be missing or not up-to-date.[/dim]")

        if cost_report_file_path:
            cls.save_to_csv(records, cost_report_file_path)

    @staticmethod
    def save_to_csv(records: list[dict[str, Any]], file_path: Path) -> None:
        """Save records to CSV file."""
        if not records:
            return

        # Collect all unique field names across all records
        all_fieldnames: set[str] = set()
        for record in records:
            all_fieldnames.update(record.keys())

        # Sort fieldnames for consistent column order
        fieldnames = sorted(all_fieldnames)

        with open(file_path, "w", newline="", encoding="utf-8") as csvfile:
            writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
            writer.writeheader()
            writer.writerows(records)

    @classmethod
    def compute_total_cost(cls, input_non_cached_cost: float, input_cached_cost: float, output_cost: float) -> float:
        return input_non_cached_cost + input_cached_cost + output_cost

    @classmethod
    def compute_cost_report(cls, llm_tokens_usage: LLMTokensUsage) -> LLMTokenCostReport:
        costs_by_token_category: CostsByCategoryDict = {}
        for token_type, nb_tokens in llm_tokens_usage.nb_tokens_by_category.items():
            cost_per_token = model_cost_per_token(
                costs=llm_tokens_usage.unit_costs,
                cost_category=token_type.to_cost_category,
            )
            costs_by_token_category[token_type.to_cost_category] = cost_per_token * nb_tokens
        return LLMTokenCostReport(
            job_metadata=llm_tokens_usage.job_metadata,
            inference_model_name=llm_tokens_usage.inference_model_name,
            platform_llm_id=llm_tokens_usage.inference_model_id,
            nb_tokens_by_category=llm_tokens_usage.nb_tokens_by_category,
            costs_by_token_category=costs_by_token_category,
        )

    @classmethod
    def complete_cost_report(cls, llm_tokens_usage: LLMTokensUsage) -> LLMTokenCostReport:
        cost_report = cls.compute_cost_report(llm_tokens_usage=llm_tokens_usage)
        # compute the input_non_cached tokens
        if cost_report.nb_tokens_by_category.get(TokenCategory.INPUT_NON_CACHED) is not None:
            msg = "CostCategory.INPUT_NON_CACHED already exists in the cost report"
            raise CostRegistryError(msg)
        # we use pop to remove input tokens which will be replaced by "input joined"
        nb_tokens_input_joined = cost_report.nb_tokens_by_category.pop(TokenCategory.INPUT, 0)
        cost_report.costs_by_token_category.pop(CostCategory.INPUT, None)

        nb_tokens_input_cached = cost_report.nb_tokens_by_category.get(TokenCategory.INPUT_CACHED, 0)
        nb_tokens_input_non_cached = nb_tokens_input_joined - nb_tokens_input_cached
        cost_report.nb_tokens_by_category[TokenCategory.INPUT_JOINED] = nb_tokens_input_joined
        cost_report.nb_tokens_by_category[TokenCategory.INPUT_NON_CACHED] = nb_tokens_input_non_cached
        cost_report.nb_tokens_by_category[TokenCategory.INPUT_CACHED] = nb_tokens_input_cached

        cost_report.costs_by_token_category[CostCategory.INPUT_NON_CACHED] = nb_tokens_input_non_cached * model_cost_per_token(
            costs=llm_tokens_usage.unit_costs,
            cost_category=CostCategory.INPUT_NON_CACHED,
        )
        costs_input_cached = cost_report.costs_by_token_category.get(CostCategory.INPUT_CACHED, 0)
        cost_report.costs_by_token_category[CostCategory.INPUT_CACHED] = costs_input_cached
        cost_report.costs_by_token_category[CostCategory.INPUT_JOINED] = (
            costs_input_cached + cost_report.costs_by_token_category[CostCategory.INPUT_NON_CACHED]
        )
        return cost_report
