from __future__ import annotations

import unittest
from pathlib import Path
from tempfile import TemporaryDirectory

import pandas as pd

from app.prediction import (
    FEATURES,
    ID_COLUMN,
    PredictionError,
    estimate_remaining_life,
    grade_info,
    interpolate_probability,
    secure_upload_name,
    validate_input_frame,
    write_prediction_workbook,
)


class DummyCurve:
    def __init__(self, x: list[float], y: list[float]) -> None:
        self.x = x
        self.y = y


class PredictionHelpersTest(unittest.TestCase):
    def test_secure_upload_name_accepts_chinese_filename(self) -> None:
        filename, suffix = secure_upload_name("管道数据.xlsx", "run123")

        self.assertEqual(suffix, ".xlsx")
        self.assertTrue(filename.endswith("_run123.xlsx"))

    def test_secure_upload_name_rejects_unsupported_extension(self) -> None:
        with self.assertRaises(PredictionError):
            secure_upload_name("管道数据.txt", "run123")

    def test_probability_helpers(self) -> None:
        self.assertEqual(interpolate_probability([1, 5, 10], [0.9, 0.8, 0.6], 6), 0.6)
        self.assertEqual(estimate_remaining_life([1, 5, 10], [0.9, 0.4, 0.2]), 5.0)

    def test_grade_boundaries(self) -> None:
        self.assertEqual(grade_info(0.2)[0], "I级")
        self.assertEqual(grade_info(0.8)[0], "IV级")
        self.assertEqual(grade_info(0.81)[0], "V级")

    def test_validate_input_frame_reports_missing_columns(self) -> None:
        df = pd.DataFrame({ID_COLUMN: [1], FEATURES[0]: [1]})

        with self.assertRaises(PredictionError) as ctx:
            validate_input_frame(df)

        self.assertIn("缺少必要字段", ctx.exception.message)

    def test_prediction_workbook_keeps_sample_data_in_one_sheet(self) -> None:
        curves = [
            DummyCurve([1, 2], [0.9, 0.7]),
            DummyCurve([1, 2], [0.8, 0.6]),
        ]
        summary_rows = [{"pipe_id": "P001"}, {"pipe_id": "P002"}]
        summary_sheet_rows = [
            {"管道编号": "P001", "健康概率": 0.7},
            {"管道编号": "P002", "健康概率": 0.6},
        ]

        with TemporaryDirectory() as temp_dir:
            output_path = Path(temp_dir) / "prediction.xlsx"
            write_prediction_workbook(output_path, curves, summary_rows, summary_sheet_rows)

            workbook = pd.ExcelFile(output_path)
            self.assertEqual(workbook.sheet_names, ["结果摘要", "样本数据"])

            sample_data = pd.read_excel(output_path, sheet_name="样本数据")
            self.assertEqual(len(sample_data), 4)
            self.assertEqual(sample_data["管道编号"].tolist(), ["P001", "P001", "P002", "P002"])
            self.assertIn("风险概率", sample_data.columns)


if __name__ == "__main__":
    unittest.main()