from __future__ import annotations

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

from pydantic import BaseModel

from srs2std.evaluate.decontextualize_context import (
    Decontextualization,
    DecontextualizeContext,
)
from srs2std.evaluate.decontextualize_strategies import Molecular
from srs2std.evaluate.factcheck_context import FactcheckContext
from srs2std.evaluate.factcheck_strategies import Thresholding, Verdict
from srs2std.evaluate.factcheck_strategies.scorers import MiniCheckRoBERTa
from srs2std.llms.gpt_oss import GPT_OSS
from srs2std.srs.instance import SRS, load_srs_md, load_srs_yaml
from srs2std.util import is_now_string, now_string


def main(
    srs_std_dirs_root: Path,
    srs_orig_dir: Path,
    srs_title: str,
    count: int,
    force: bool,
):
    decontextualize_srs(srs_orig_dir, srs_title, count, force)

    srs_std_dirs = list_generated_paths(srs_std_dirs_root)
    for srs_std_dir in srs_std_dirs:
        decontextualize_srs(srs_std_dir, srs_title, count, force)

        # precision
        factcheck_results = factcheck(srs_orig_dir, srs_std_dir)
        save_evaluation(srs_std_dir, "precision", factcheck_results)

        # recall
        factcheck_results = factcheck(srs_std_dir, srs_orig_dir)
        save_evaluation(srs_std_dir, "recall", factcheck_results)


###############################################################################


def decontextualize_srs(
    srs_dir: Path,
    srs_title: str,
    count: int,
    force: bool,
):
    decontextualization_dir = srs_dir / "decontextualization"
    decontextualization_dir.mkdir(parents=True, exist_ok=True)

    if force:
        delete_generated_paths(decontextualization_dir)
    decontextualization_paths = list_generated_paths(decontextualization_dir)
    if len(decontextualization_paths) >= count:
        return

    for _ in range(count - len(decontextualization_paths)):
        srs = load_srs_from_dir(srs_dir)

        context = DecontextualizeContext(srs_title=srs_title, srs=srs)
        strategy = Molecular(llm=GPT_OSS())
        results = context.run(strategy)

        save_decontextualizations(decontextualization_dir, results)


def delete_generated_paths(root_dir: Path):
    for path in list_generated_paths(root_dir):
        if path.is_file():
            path.unlink()


def list_generated_paths(root_dir: Path) -> List[Path]:
    if not root_dir.is_dir():
        return []
    return [p for p in root_dir.iterdir() if is_now_string(p.stem)]


def load_srs_from_dir(srs_dir: Path) -> SRS:
    # first, search post-processed version in .yaml
    yaml_path = srs_dir / "postprocessed.yaml"
    if yaml_path.is_file():
        return load_srs_yaml(yaml_path)

    # second, search pre-processed version in .md
    md_path = srs_dir / "preprocessed.md"
    return load_srs_md(md_path)


def load_decontextualizations(path: Path) -> List[Decontextualization]:
    with path.open(encoding="utf-8") as fp:
        data = json.load(fp)
    return [Decontextualization(**item) for item in data]


def save_decontextualizations(out_dir: Path, items: List[Decontextualization]):
    path = out_dir / f"{now_string()}.json"
    data = [json.loads(item.model_dump_json()) for item in items]
    with path.open("w", encoding="utf-8") as fp:
        json.dump(data, fp, indent=4, ensure_ascii=False)


###############################################################################


class FactcheckResult(BaseModel):
    decontextualization_path: Path
    verdicts: List[Verdict]


def factcheck(
    srs_ground_dir: Path,
    srs_check_dir: Path,
) -> List[FactcheckResult]:
    srs_ground = load_srs_from_dir(srs_ground_dir)

    decontextualization_dir = srs_check_dir / "decontextualization"
    decontextualization_paths = list_generated_paths(decontextualization_dir)

    factcheck_scorer = MiniCheckRoBERTa()
    factcheck_strategy = Thresholding(threshold=0.5, scorer=factcheck_scorer)

    factcheck_results = []
    for decontextualization_path in decontextualization_paths:
        decontextualizations = load_decontextualizations(decontextualization_path)

        ground_doc = srs_ground.to_str(ignore_title=False, ignore_headers=False)
        check_sents = [dectx.after for dectx in decontextualizations]
        factcheck_context = FactcheckContext(ground_doc, check_sents)
        verdicts = factcheck_context.run(factcheck_strategy)

        factcheck_result = FactcheckResult(
            decontextualization_path=decontextualization_path,
            verdicts=verdicts,
        )
        factcheck_results.append(factcheck_result)

    return factcheck_results


###############################################################################


class EvaluationResult(BaseModel):
    decontextualization_path: Path
    verdicts: List[Verdict]

    def to_summary_json(self):
        labels = [verdict.conclusion for verdict in self.verdicts]
        num_of_sentences = len(labels)
        num_of_supported = sum(labels)
        return {
            "decontextualization_path": str(self.decontextualization_path),
            "num_of_sentences": num_of_sentences,
            "num_of_supported": num_of_supported,
        }


def save_evaluation(
    srs_std_dir: Path,
    metric_name: Literal["precision", "recall"],
    factcheck_results: List[FactcheckResult],
):
    output_dir = srs_std_dir / "evaluation" / metric_name
    output_dir.mkdir(parents=True, exist_ok=True)

    evaluation_results = []
    for factcheck_result in factcheck_results:
        evaluation_result = EvaluationResult(
            decontextualization_path=factcheck_result.decontextualization_path,
            verdicts=factcheck_result.verdicts,
        )
        evaluation_results.append(evaluation_result)

    save_verdicts(output_dir, evaluation_results)
    save_summary(output_dir, evaluation_results)


def save_verdicts(output_dir: Path, evaluation_results: List[EvaluationResult]):
    for evaluation_result in evaluation_results:
        verdicts_json = [
            json.loads(verdict.model_dump_json())
            for verdict in evaluation_result.verdicts
        ]
        filename = f"{evaluation_result.decontextualization_path.stem}.json"
        output_path = output_dir / filename
        with output_path.open("w", encoding="utf-8") as fp:
            json.dump(verdicts_json, fp, indent=4, ensure_ascii=False)


def save_summary(output_dir: Path, evaluation_results: List[EvaluationResult]):
    summary_json = [
        evaluation_result.to_summary_json() for evaluation_result in evaluation_results
    ]
    filename = "summary.json"
    output_path = output_dir / filename
    with output_path.open("w", encoding="utf-8") as fp:
        json.dump(summary_json, fp, indent=4, ensure_ascii=False)
