from __future__ import annotations

import hashlib
import json
import uuid
from dataclasses import asdict, dataclass, field, is_dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any


PROJECT_MODEL_VERSION = "1.0"
DEFAULT_CASES = ("existing", "proposed")
DEFAULT_FREQUENCIES = (1.0, 2.0, 5.0, 10.0)


@dataclass
class RainStation:
    id: str
    name: str = ""
    x: float | None = None
    y: float | None = None
    elevation: float | None = None
    weight: float | None = None


@dataclass
class BasinData:
    area: float | None = None
    mainLength: float | None = None
    mainSlope: float | None = None
    rainStations: list[RainStation] = field(default_factory=list)
    soil: str = ""
    vegetation: str = ""


@dataclass
class SectionPoint:
    offset: float
    elevation: float


@dataclass
class SectionData:
    id: str
    chainage: float | None = None
    geometry: list[SectionPoint] = field(default_factory=list)
    roughness: float | None = None
    leftBankElev: float | None = None
    rightBankElev: float | None = None
    bedElev: float | None = None


@dataclass
class StructureData:
    id: str
    structureType: str = ""
    piers: list[dict[str, Any]] = field(default_factory=list)
    occupyArea: float | None = None
    spanLength: float | None = None
    pierCount: int | None = None
    pierWidth: float | None = None
    naturalArea: float | None = None
    effectiveArea: float | None = None
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class ResultRecord:
    resultId: str
    programId: str
    sourceLink: str
    parameterVersion: str
    version: str
    createdAt: str
    inputs: dict[str, Any] = field(default_factory=dict)
    outputs: dict[str, Any] = field(default_factory=dict)
    dependencyVersions: dict[str, str] = field(default_factory=dict)
    calcBookPath: str = ""


@dataclass
class FrequencyData:
    P: float
    returnPeriod: float | None = None
    results: dict[str, ResultRecord] = field(default_factory=dict)


@dataclass
class CaseData:
    id: str
    name: str
    caseType: str
    frequencies: dict[str, FrequencyData] = field(default_factory=dict)
    structures: list[StructureData] = field(default_factory=list)


@dataclass
class ProjectData:
    id: str
    name: str
    type: str
    basin: BasinData = field(default_factory=BasinData)
    sections: list[SectionData] = field(default_factory=list)
    cases: dict[str, CaseData] = field(default_factory=dict)
    metadata: dict[str, Any] = field(default_factory=dict)
    schemaVersion: str = PROJECT_MODEL_VERSION


# @hook B3 20260804
# Shared project schema and JSON persistence for basin/sections/cases/frequencies/structures/results.
def new_project(project_id: str | None = None, name: str = "", project_type: str = "洪水影响评价") -> ProjectData:
    project = ProjectData(
        id=project_id or f"PRJ-{uuid.uuid4().hex[:12]}",
        name=name,
        type=project_type,
        metadata={"createdAt": _now_iso()},
    )
    project.cases["existing"] = CaseData(id="existing", name="现状", caseType="existing")
    project.cases["proposed"] = CaseData(id="proposed", name="工程后", caseType="proposed")
    for case in project.cases.values():
        for frequency in DEFAULT_FREQUENCIES:
            case.frequencies[_frequency_key(frequency)] = FrequencyData(
                P=frequency,
                returnPeriod=_return_period(frequency),
            )
    return project


def save_project(project: ProjectData, path: str | Path) -> Path:
    target = Path(path)
    target.parent.mkdir(parents=True, exist_ok=True)
    target.write_text(
        json.dumps(to_dict(project), ensure_ascii=False, indent=2, sort_keys=True),
        encoding="utf-8",
    )
    return target


def load_project(path: str | Path) -> ProjectData:
    payload = json.loads(Path(path).read_text(encoding="utf-8"))
    return project_from_dict(payload)


def to_dict(value: Any) -> Any:
    if is_dataclass(value):
        return {key: to_dict(item) for key, item in asdict(value).items()}
    if isinstance(value, dict):
        return {str(key): to_dict(item) for key, item in value.items()}
    if isinstance(value, list):
        return [to_dict(item) for item in value]
    return value


def project_from_dict(payload: dict[str, Any]) -> ProjectData:
    basin_payload = payload.get("basin") or {}
    basin = BasinData(
        area=basin_payload.get("area"),
        mainLength=basin_payload.get("mainLength"),
        mainSlope=basin_payload.get("mainSlope"),
        rainStations=[RainStation(**row) for row in basin_payload.get("rainStations", [])],
        soil=basin_payload.get("soil", ""),
        vegetation=basin_payload.get("vegetation", ""),
    )
    sections = [
        SectionData(
            id=row["id"],
            chainage=row.get("chainage"),
            geometry=[SectionPoint(**point) for point in row.get("geometry", [])],
            roughness=row.get("roughness"),
            leftBankElev=row.get("leftBankElev"),
            rightBankElev=row.get("rightBankElev"),
            bedElev=row.get("bedElev"),
        )
        for row in payload.get("sections", [])
    ]
    cases = {
        case_id: _case_from_dict(case_payload)
        for case_id, case_payload in (payload.get("cases") or {}).items()
    }
    return ProjectData(
        id=payload["id"],
        name=payload.get("name", ""),
        type=payload.get("type", "洪水影响评价"),
        basin=basin,
        sections=sections,
        cases=cases,
        metadata=payload.get("metadata", {}),
        schemaVersion=payload.get("schemaVersion", PROJECT_MODEL_VERSION),
    )


def add_result(
    project: ProjectData,
    case_id: str,
    frequency: float | str,
    program_id: str,
    inputs: dict[str, Any],
    outputs: dict[str, Any],
    *,
    source_link: str | None = None,
    result_version: str = "1",
    calc_book_path: str = "",
    dependency_versions: dict[str, str] | None = None,
) -> ResultRecord:
    frequency_data = get_frequency(project, case_id, frequency, create=True)
    canonical_inputs = canonicalize_program_fields(program_id, inputs)
    canonical_outputs = canonicalize_program_fields(program_id, outputs)
    if dependency_versions is None:
        dependency_versions = dependency_versions_for(project, case_id, frequency, program_id)
    record = ResultRecord(
        resultId=f"RES-{uuid.uuid4().hex[:12]}",
        programId=program_id,
        sourceLink=source_link or f"program_registry:{program_id}",
        parameterVersion=make_parameter_version(canonical_inputs),
        version=result_version,
        createdAt=_now_iso(),
        inputs=canonical_inputs,
        outputs=canonical_outputs,
        dependencyVersions=dict(dependency_versions),
        calcBookPath=calc_book_path,
    )
    frequency_data.results[program_id] = record
    return record


# @hook B5 20260804
# Tracks upstream parameter versions so downstream results can be marked stale after upstream recalculation.
def get_completed_programs(project: ProjectData, case_id: str = "proposed", frequency: float | str = 1.0) -> set[str]:
    try:
        frequency_data = get_frequency(project, case_id, frequency)
    except KeyError:
        return set()
    return set(frequency_data.results)


def dependency_versions_for(
    project: ProjectData,
    case_id: str,
    frequency: float | str,
    program_id: str,
) -> dict[str, str]:
    try:
        from .program_registry import get_program

        dependencies = get_program(program_id).dependencies
    except Exception:  # noqa: BLE001 - project files must remain loadable if registry changes.
        dependencies = ()
    try:
        frequency_data = get_frequency(project, case_id, frequency)
    except KeyError:
        return {}
    versions: dict[str, str] = {}
    for dependency in dependencies:
        record = frequency_data.results.get(dependency)
        if record is not None:
            versions[dependency] = record.parameterVersion
    return versions


def list_stale_results(
    project: ProjectData,
    case_id: str | None = None,
    frequency: float | str | None = None,
) -> list[dict[str, Any]]:
    stale: list[dict[str, Any]] = []
    case_items = (
        [(case_id, get_case(project, case_id))]
        if case_id is not None and case_id in project.cases
        else list(project.cases.items())
    )
    for current_case_id, case in case_items:
        frequency_items = (
            [(_frequency_key(frequency), get_frequency(project, current_case_id, frequency))]
            if frequency is not None and _frequency_key(frequency) in case.frequencies
            else list(case.frequencies.items())
        )
        for _frequency_id, frequency_data in frequency_items:
            for program_id, record in frequency_data.results.items():
                current_versions = dependency_versions_for(project, current_case_id, frequency_data.P, program_id)
                changed: list[dict[str, str]] = []
                for dependency, current_version in current_versions.items():
                    stored_version = record.dependencyVersions.get(dependency, "")
                    if stored_version and stored_version != current_version:
                        changed.append(
                            {
                                "programId": dependency,
                                "storedVersion": stored_version,
                                "currentVersion": current_version,
                            }
                        )
                if changed:
                    stale.append(
                        {
                            "caseId": current_case_id,
                            "frequency": frequency_data.P,
                            "programId": program_id,
                            "message": "数据过期，需重算",
                            "staleDependencies": changed,
                        }
                    )
    return stale


def get_frequency(project: ProjectData, case_id: str, frequency: float | str, *, create: bool = False) -> FrequencyData:
    case = get_case(project, case_id, create=create)
    key = _frequency_key(frequency)
    if key not in case.frequencies:
        if not create:
            raise KeyError(f"frequency not found: {frequency}")
        p_value = float(frequency)
        case.frequencies[key] = FrequencyData(P=p_value, returnPeriod=_return_period(p_value))
    return case.frequencies[key]


def get_case(project: ProjectData, case_id: str, *, create: bool = False) -> CaseData:
    if case_id not in project.cases:
        if not create:
            raise KeyError(f"case not found: {case_id}")
        project.cases[case_id] = CaseData(id=case_id, name=case_id, caseType=case_id)
    return project.cases[case_id]


def make_parameter_version(params: dict[str, Any]) -> str:
    encoded = json.dumps(params, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode("utf-8")
    return hashlib.sha256(encoded).hexdigest()[:16]


def build_schema_example() -> dict[str, Any]:
    project = new_project("PRJ-DEMO", "桥梁洪评示例工程", "桥梁洪水影响评价")
    project.basin = BasinData(area=47.02, mainLength=12.4, mainSlope=9.8, soil="中等土", vegetation="草灌")
    project.sections.append(
        SectionData(
            id="SEC-001",
            chainage=0.0,
            geometry=[SectionPoint(0.0, 1250.0), SectionPoint(20.0, 1248.6), SectionPoint(45.0, 1250.2)],
            roughness=0.035,
            bedElev=1248.6,
        )
    )
    project.cases["proposed"].structures.append(
        StructureData(
            id="BR-001",
            structureType="桥梁",
            piers=[{"pierWidth": 1.5, "count": 2}],
            occupyArea=12.0,
            spanLength=30.0,
            pierCount=2,
            pierWidth=1.5,
        )
    )
    add_result(
        project,
        "proposed",
        1.0,
        "3-17",
        {
            "design_water_level_m": 1252.41,
            "backwater_height_m": 0.18,
            "safety_freeboard_m": 0.5,
            "girder_bottom_elevation_m": 1253.2,
        },
        {"clearance_margin_m": 0.11, "pass": True},
    )
    return to_dict(project)


STR_FIELD_ALIASES: dict[str, dict[str, str]] = {
    "3-15": {
        "natural_area_m2": "naturalArea",
        "effective_area_m2": "effectiveArea",
        "design_flow_m3s": "peakQ",
        "natural_width_m": "naturalWidth",
        "effective_width_m": "effectiveWidth",
        "blocked_flow_ratio_percent": "blockedFlowRatio",
        "obstruction_area_ratio_percent": "blockRatio",
    },
    "3-16": {
        "design_flow_m3s": "peakQ",
        "full_section_area_m2": "naturalArea",
        "bridge_actual_area_m2": "effectiveArea",
        "natural_channel_velocity_ms": "velocity",
        "natural_water_surface_slope": "waterSurfaceSlope",
        "obstruction_area_ratio_percent": "blockRatio",
        "blocked_flow_ratio_percent": "blockedFlowRatio",
        "flow_angle_deg": "flowAngle",
        "design_water_level_m": "desWL",
        "max_backwater_height_m": "backwater",
        "backwater_curve_length_m": "backwaterLength",
    },
    "3-13": {
        "bed_elevation_m": "bedElev",
        "general_scour_depth_m": "generalScourDepth",
        "local_scour_depth_m": "localScourDepth",
        "scour_line_elevation_m": "scourLineElev",
        "Qp": "peakQ",
        "v": "velocity",
        "hp": "hydraulicRadius",
        "B1": "pierWidth",
    },
    "3-17": {
        "design_water_level_m": "desWL",
        "backwater_height_m": "backwater",
        "safety_freeboard_m": "safetyFreeboard",
        "girder_bottom_elevation_m": "girderBottomElev",
        "clearance_margin_m": "clearanceMargin",
    },
    "3-12": {
        "bed_elevation_m": "bedElev",
        "general_scour_depth_m": "generalScourDepth",
        "local_scour_depth_m": "localScourDepth",
        "foundation_bottom_elevation_m": "foundationBottomElev",
        "required_embedment_m": "requiredBurialDepth",
        "embedment_margin_m": "burialDepthMargin",
    },
    "3-14": {
        "max_backwater_height_m": "backwater",
        "backwater_length_m": "backwaterLength",
        "design_water_level_m": "desWL",
        "bed_slope": "bedSlope",
        "curve_points": "backwaterCurve",
    },
}


def canonicalize_program_fields(program_id: str, params: dict[str, Any]) -> dict[str, Any]:
    aliases = STR_FIELD_ALIASES.get(program_id, {})
    canonical: dict[str, Any] = {}
    for key, value in params.items():
        canonical[aliases.get(key, key)] = value
    if program_id == "3-12" and "burialDepthMargin" in canonical and "burialDepth" not in canonical:
        canonical["burialDepth"] = canonical["burialDepthMargin"]
    return canonical


def _case_from_dict(payload: dict[str, Any]) -> CaseData:
    return CaseData(
        id=payload["id"],
        name=payload.get("name", payload["id"]),
        caseType=payload.get("caseType", payload["id"]),
        frequencies={
            key: _frequency_from_dict(row)
            for key, row in (payload.get("frequencies") or {}).items()
        },
        structures=[
            StructureData(
                id=row["id"],
                structureType=row.get("structureType", ""),
                piers=row.get("piers", []),
                occupyArea=row.get("occupyArea"),
                spanLength=row.get("spanLength"),
                pierCount=row.get("pierCount"),
                pierWidth=row.get("pierWidth"),
                naturalArea=row.get("naturalArea"),
                effectiveArea=row.get("effectiveArea"),
                metadata=row.get("metadata", {}),
            )
            for row in payload.get("structures", [])
        ],
    )


def _frequency_from_dict(payload: dict[str, Any]) -> FrequencyData:
    return FrequencyData(
        P=payload["P"],
        returnPeriod=payload.get("returnPeriod"),
        results={
            key: ResultRecord(**row)
            for key, row in (payload.get("results") or {}).items()
        },
    )


def _frequency_key(frequency: float | str) -> str:
    value = float(frequency)
    if value.is_integer():
        return str(int(value))
    return str(value).rstrip("0").rstrip(".")


def _return_period(frequency: float) -> float | None:
    if frequency <= 0:
        return None
    return round(100.0 / frequency, 6)


def _now_iso() -> str:
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat()
