import json
import logging
from pathlib import Path
from typing import Dict, List, Literal

import pandas as pd
from huggingface_hub import hf_hub_download
from inspect_ai import Task, task
from inspect_ai.dataset import MemoryDataset, Sample
from inspect_ai.model import GenerateConfig, Model, get_model
from inspect_ai.scorer import Score, Scorer, Target, mean, scorer, stderr
from inspect_ai.solver import Generate, Solver, TaskState, solver, use_tools
from inspect_ai.tool import ToolDef
from pydantic import BaseModel, ValidationError

from astabench.constants import ASTA_BENCH_DATASET_REPO, ASTA_BENCH_DATASET_REVISION
from astabench.evals.arxivdigestables.prompts import (
    eval_prompt,
    input_prompt,
    unroll_prompt,
)
from astabench.evals.utils import not_implemented_solver
from astabench.tools import make_asta_s2_mcp_tools, table_editor

logger = logging.getLogger(__name__)

# TODO: Add a hard version!
"""
Pydantic models defining expected output format for table generation.
"""


class Cell(BaseModel):
    """
    A Cell Object consists of a paper ID, a column name and
    the corresponding cell value at that row & column in the table.
    """

    paper_id: str
    column_name: str
    cell_value: str


class Table(BaseModel):
    """
    A Table Object is a List of Cell Objects.
    """

    cell_values: List[Cell]


def convert_pydantic_table_to_json(table: Table) -> dict:
    """
    Convert a generated table (in our Pydantic format)
    to the json format expected by unrolling/eval code.
    """
    table_json: dict[str, dict[str, str]] = {}
    for cell in table.cell_values:
        if cell.column_name not in table_json:
            table_json[cell.column_name] = {}
        table_json[cell.column_name][cell.paper_id] = cell.cell_value
    return table_json


@scorer(metrics=[mean(), stderr()])
def score_tables(unroller_model: str | Model, scorer_model: str | Model) -> Scorer:
    unroller_model = get_model(unroller_model)
    scorer_model = get_model(scorer_model)

    async def score(state: TaskState, target: Target) -> Score:
        # Validate that output is formatted correctly

        try:
            tjson = state.output.completion
            tjson = tjson[tjson.index("{") : tjson.rindex("}") + 1]
        except ValueError as e:
            msg = str(e)
            if "substring not found" in msg:
                msg = "JSON object not found (missing `{` or `}`)"
            return Score(
                value=0.0,
                answer=state.output.completion,
                explanation=f"Error parsing response: {msg}",
            )

        try:
            table = Table.model_validate_json(tjson)
        except ValidationError as ex:
            # If output is not json-formatted correctly, return a score of 0
            return Score(
                value=0.0,
                answer=state.output.completion,
                explanation=f"Error parsing response: {ex}",
            )

        # If yes, then unroll the generated table for evaluation
        table_md = pd.DataFrame(convert_pydantic_table_to_json(table)).to_markdown()
        final_input = unroll_prompt.replace("[TABLE]", table_md)
        unrolled_table_str = await unroller_model.generate(
            final_input,
            config=GenerateConfig(
                temperature=0.0, max_connections=10, response_schema=None
            ),
        )
        # Get individual statements from the unrolled table
        unrolled_table = unrolled_table_str.completion.strip()
        unrolled_target_table = [
            line for line in target.text.split("\n") if line.strip()
        ]
        # Compare each statement against unrolled target statements
        # and use that to compute the recall for the generated table
        final_input = eval_prompt.replace("[UNROLLED_TABLE]", unrolled_table)
        responses = []
        recall = 0.0
        for statement in unrolled_target_table:
            response = await scorer_model.generate(
                final_input.replace("[STATEMENT]", statement),
                config=GenerateConfig(
                    temperature=0.0, max_connections=10, response_schema=None
                ),
            )
            supported_stmt = response.completion.lower()
            responses.append({"statement": statement, "supported": supported_stmt})
            if "true" in supported_stmt:
                recall += 1
        target_statement_num = len(unrolled_target_table)
        recall /= target_statement_num
        return Score(
            value=recall,
            answer=state.output.completion,
            explanation=json.dumps(responses),
        )

    return score


def json_to_sample(obj: dict[str, str]) -> Sample:
    """Convert arxivdigestables instance to inspect_ai Sample format."""
    return Sample(id=obj["id"], input=obj["input"], target=obj["target"])


def make_paper_list_input(paper_list: List[Dict]) -> str:
    """Format list of papers so that it can be used in the prompt."""
    paper_text = ""
    for index, paper in enumerate(paper_list):
        abstract = (
            paper["abstract"].strip()
            if "abstract" in paper and paper["abstract"]
            else None
        )
        title = paper["title"]
        corpus_id = paper["corpus_id"]
        paper_text += f"Paper {corpus_id} title: {title}\nPaper {corpus_id} abstract: {abstract}\n\n"

    return paper_text


@solver
def setup_snippet_tool() -> Solver:
    # We set an insertion date so that the publication date arg gets pruned
    # from the tool (no reason for the model to see it); it doesn't really
    # matter since we are also forcing the corpus ids, which are all older than
    # 2025 for this task
    base_tools = make_asta_s2_mcp_tools(insertion_date="2026-01-01")
    tools_by_name = {ToolDef(tool).name: tool for tool in base_tools}

    if "snippet_search" not in tools_by_name:
        raise KeyError(
            "The 'snippet_search' tool is not available in the S2 MCP tools; this is unexpected (has the S2 MCP server changed?)"
        )
    snippet_tool = tools_by_name["snippet_search"]

    async def solve(state: TaskState, generate: Generate) -> TaskState:

        async def inner_snippet_search(
            query: str,
            limit: int,
            paper_ids: str | None = None,
        ) -> str:
            allowed_paper_id_list = [
                f"CorpusId:{x}" for x in state.metadata["corpus_ids"]
            ]
            paper_id_list = allowed_paper_id_list
            if paper_ids:
                paper_ids = [
                    x for x in paper_ids.split(",") if x.startswith("CorpusId:")
                ]
                # Filter out any corpus ids that aren't part of the allowed
                # list; currently we don't return any warning about filtered
                # ones unless all of them get filtered.
                paper_id_list = list(
                    set(allowed_paper_id_list) & set(map(str, paper_ids))
                )
                if not paper_id_list:
                    return json.dumps(
                        {
                            "error": "No valid corpus ids found in the input.",
                            "allowed_paper_ids": allowed_paper_id_list,
                        }
                    )
            try:
                return str(await snippet_tool(query, limit, paper_ids=paper_id_list))
            except Exception as e:
                return json.dumps({"exception": str(e)})

        restricted_snippet_td = ToolDef(snippet_tool)
        restricted_snippet_td.tool = inner_snippet_search
        return await use_tools(restricted_snippet_td.as_tool())(state, generate)

    return solve


@task
def arxivdigestables(
    unroller_model: str | Model = "openai/gpt-4o",
    scorer_model: str | Model = "openai/gpt-4o",
    with_snippet_search_tool: bool = True,
    with_table_editor_tool: bool = False,
    split: Literal["test", "val"] = "test",
) -> Task:
    """
    ArxivDIGESTables task, using paper titles and abstracts
    as input to produce comparison tables.
    """
    # Load arxivDIGESTables dataset from Huggingface.
    # We use a manually curated and filtered hard version for testing.
    hf_repo_id = ASTA_BENCH_DATASET_REPO

    hf_file_path = f"tasks/arxivdigestables/final_{split}_set.jsonl"
    cached_file = hf_hub_download(
        repo_id=hf_repo_id,
        filename=hf_file_path,
        repo_type="dataset",
        revision=ASTA_BENCH_DATASET_REVISION,
    )
    instances = [
        json.loads(line)
        for line in Path(cached_file).read_text(encoding="utf-8").splitlines()
    ]

    # Load unrolled tables (for evaluation) from json file
    hf_file_path = "tasks/arxivdigestables/unrolled_tables.json"
    cached_file = hf_hub_download(
        repo_id=hf_repo_id,
        filename=hf_file_path,
        repo_type="dataset",
        revision=ASTA_BENCH_DATASET_REVISION,
    )
    unrolled_tables = json.loads(Path(cached_file).read_text(encoding="utf-8"))

    formatted_dataset = []
    for instance in instances:
        instance_id = instance["tabid"]
        if instance_id not in unrolled_tables:
            logger.warning(
                f"Skipping instance {instance_id} because it has no unrolled table."
            )
            continue
        # Get number of columns and rows in gold table to add to input prompt.
        table = instance["table"]
        instance_col_num = len(table)
        first_col = list(table.keys())[0]
        instance_paper_num = len(table[first_col])
        # Add the number of columns and rows to the input prompt.
        final_input = input_prompt.replace("[COL_NUM]", str(instance_col_num))
        final_input = final_input.replace("[PAPER_NUM]", str(instance_paper_num))
        # Add table caption to the input prompt.
        final_input = final_input.replace("[CAPTION]", instance["caption"])
        # Format paper titles/abstracts and add to the input prompt.
        paper_input = make_paper_list_input(instance["row_bib_map"])
        final_input += paper_input

        final_input += "\n\nRespond with the following json schema:\n" + json.dumps(
            Table.schema(), indent=2
        )
        formatted_dataset.append(
            Sample(
                id=instance_id,
                input=final_input,
                target=unrolled_tables[instance_id],
                metadata={
                    "corpus_ids": [str(x["corpus_id"]) for x in instance["row_bib_map"]]
                },
            )
        )
    tool_setups = []
    if with_snippet_search_tool:
        tool_setups.append(setup_snippet_tool())
    if with_table_editor_tool:
        tool_setups.append(use_tools(table_editor()))
    return Task(
        dataset=MemoryDataset(formatted_dataset),
        setup=tool_setups or None,
        solver=not_implemented_solver(),
        scorer=score_tables(unroller_model, scorer_model),
    )


@task
def arxivdigestables_test(
    *args,
    **kwargs,
) -> Task:
    """ArxivDIGESTables eval (test split)."""
    kwargs.pop("split", None)
    return arxivdigestables(
        *args,
        **kwargs,
        split="test",
    )


@task
def arxivdigestables_validation(
    *args,
    **kwargs,
) -> Task:
    """ArxivDIGESTables eval (validation split)."""
    kwargs.pop("split", None)
    return arxivdigestables(
        *args,
        **kwargs,
        split="val",
    )


__all__ = [
    "arxivdigestables",
    "arxivdigestables_test",
    "arxivdigestables_validation",
]
