#!/usr/bin/env python3
from __future__ import annotations

import time
from typing import Any

from .io import (
    MYCHEM_CACHE_PATH,
    SOURCE_CACHE_FILES,
    WHO_API_URL,
    WHO_DEVICES_URL,
    WHO_FINISHED_URL,
    WHO_VACCINES_URL,
    clean_text,
    diff_payloads,
    ensure_work_dirs,
    iso_now,
    read_json,
    round_seconds,
    sha256_json,
    source_snapshot,
    time_call,
    write_json,
)
from .loaders import load_apis, load_vaccines
from .probes import PROBE_RESULT_PATHS, PROBE_SUMMARY_PATH, run_probe_suite
from .types import WhoFullScaleArtifacts

REGRESSION_CONTROL_PATH = PROBE_SUMMARY_PATH.parent / "who_regression_control.json"
VALIDATION_PATH = PROBE_SUMMARY_PATH.parent / "who_validation.json"
DESIGN_PATH = PROBE_SUMMARY_PATH.parent / "who_loader_design.json"
SAMPLES_PATH = PROBE_SUMMARY_PATH.parent / "who_sample_records.json"
FULL_SCALE_RESULTS_PATH = PROBE_SUMMARY_PATH.parent / "who_full_scale_results.json"

VALIDATION_TERMS = {
    "BCG": ["bcg"],
    "measles": ["measles"],
    "HPV": ["human papillomavirus", "hpv"],
    "COVID-19": ["covid-19"],
    "yellow fever": ["yellow fever"],
}

MEANINGFUL_VACCINE_FIELDS = [
    "Vaccine Type",
    "Commercial Name",
    "Presentation",
    "Manufacturer",
    "Responsible NRA",
    "Date of Prequalification",
]


def current_cache_state() -> dict[str, Any]:
    return {
        "upstream_sources_cached": {
            path.name: path.is_file() for path in sorted(SOURCE_CACHE_FILES.values())
        },
        "mychem_cache_exists": MYCHEM_CACHE_PATH.is_file(),
        "mychem_cache_path": str(MYCHEM_CACHE_PATH.resolve()),
    }


def build_probe_projection(name: str, payload: dict[str, Any]) -> dict[str, Any]:
    if name == "schema_comparison":
        direct = payload["current_contract_direct_header_coverage"]
        return {
            "record_counts": payload["record_counts"],
            "direct_required_header_coverage": {
                key: {
                    "shared": value["shared"],
                    "shared_count": value["shared_count"],
                }
                for key, value in direct.items()
            },
        }
    if name == "vaccine_identity":
        strategy_projection = {}
        for key, value in payload["strategy_results"].items():
            metrics = [
                "rows",
                "no_query_rows",
                "exact_match_rows",
                "exact_match_rate",
                "phrase_or_exact_rows",
                "phrase_or_exact_rate",
                "any_hit_rows",
                "any_hit_rate",
                "all_components_exact_rows",
                "all_components_exact_rate",
                "all_components_phrase_or_exact_rows",
                "all_components_phrase_or_exact_rate",
                "all_components_any_hit_rows",
                "all_components_any_hit_rate",
                "mean_component_phrase_or_exact_coverage_rate",
                "no_component_rows",
            ]
            strategy_projection[key] = {
                metric: value[metric] for metric in metrics if metric in value
            }
        return {
            "row_count": payload["row_count"],
            "strategy_results": strategy_projection,
            "winner": payload["winner"],
        }
    if name == "api_linkage":
        return {
            "row_count": payload["row_count"],
            "mychem_identity_results": {
                key: {
                    metric: value[metric]
                    for metric in [
                        "rows",
                        "no_query_rows",
                        "exact_match_rows",
                        "exact_match_rate",
                        "phrase_or_exact_rows",
                        "phrase_or_exact_rate",
                        "any_hit_rows",
                        "any_hit_rate",
                    ]
                }
                for key, value in payload["mychem_identity_results"].items()
            },
            "finished_pharma_overlap": {
                key: payload["finished_pharma_overlap"][key]
                for key in [
                    "rows",
                    "exact_inn_link_rows",
                    "exact_inn_link_rate",
                    "component_link_rows",
                    "component_link_rate",
                ]
            },
        }
    if name == "metadata_and_devices":
        return {
            "vaccine_row_count": payload["vaccine_row_count"],
            "vaccine_field_completeness": {
                key: value["percent"] for key, value in payload["vaccine_field_completeness"].items()
            },
            "vaccine_specific_field_assessment": {
                key: value["completeness"]["percent"]
                for key, value in payload["vaccine_specific_field_assessment"].items()
            },
            "ticket_validation_samples": {
                key: value["count"] for key, value in payload["ticket_validation_samples"].items()
            },
            "immunization_device_catalog": {
                "category_count": payload["immunization_device_catalog"]["category_count"],
                "item_count": payload["immunization_device_catalog"]["item_count"],
            },
        }
    raise KeyError(f"Unknown probe name: {name}")


def build_validation_payload() -> dict[str, Any]:
    vaccines = load_vaccines().rows
    probes: dict[str, Any] = {}
    all_passed = True
    for label, terms in VALIDATION_TERMS.items():
        matches = []
        for row in vaccines:
            haystack = " | ".join(
                [row.get("Vaccine Type", ""), row.get("Commercial Name", "")]
            ).lower()
            if any(term in haystack for term in terms):
                matches.append(row)
        metadata_rich = []
        for row in matches:
            present_fields = [
                field for field in MEANINGFUL_VACCINE_FIELDS if clean_text(row.get(field))
            ]
            if len(present_fields) == len(MEANINGFUL_VACCINE_FIELDS):
                metadata_rich.append(row)
        passed = bool(matches) and bool(metadata_rich)
        all_passed = all_passed and passed
        probes[label] = {
            "count": len(matches),
            "rows_with_meaningful_metadata": len(metadata_rich),
            "meaningful_metadata_fields": MEANINGFUL_VACCINE_FIELDS,
            "passed": passed,
            "samples": metadata_rich[:2] or matches[:2],
        }
    return {
        "generated_at": iso_now(),
        "overall_passed": all_passed,
        "validation_definition": {
            "required_fields": MEANINGFUL_VACCINE_FIELDS,
            "description": "Each target vaccine must appear in the WHO export and at least one matched row must carry all required metadata fields.",
        },
        "probes": probes,
    }


def write_validation_result() -> dict[str, Any]:
    payload = build_validation_payload()
    write_json(VALIDATION_PATH, payload)
    return payload


def build_sample_records_payload(
    probe_payloads: dict[str, dict[str, Any]],
    validation_payload: dict[str, Any],
) -> dict[str, Any]:
    vaccine_rows = load_vaccines().rows
    api_rows = load_apis().rows
    vaccine_samples = {
        label: details["samples"] for label, details in validation_payload["probes"].items()
    }
    api_rows_by_id = {row["WHO Product ID"]: row for row in api_rows}

    def enrich_api_samples(items: list[dict[str, Any]]) -> list[dict[str, Any]]:
        enriched = []
        for item in items:
            base_row = api_rows_by_id.get(item["who_product_id"], {})
            enriched.append(
                {
                    "who_product_id": item["who_product_id"],
                    "source_row": base_row,
                    "analysis": item,
                }
            )
        return enriched

    api_payload = probe_payloads["api_linkage"]
    return {
        "generated_at": iso_now(),
        "vaccines": {
            "validation_anchor_samples": vaccine_samples,
            "first_rows": vaccine_rows[:5],
        },
        "apis": {
            "exact_overlap_examples": enrich_api_samples(
                api_payload["finished_pharma_overlap"]["sample_exact_links"]
            ),
            "component_overlap_examples": enrich_api_samples(
                api_payload["finished_pharma_overlap"]["sample_component_links"]
            ),
            "finished_pharma_unresolved_examples": enrich_api_samples(
                api_payload["finished_pharma_overlap"]["sample_unresolved_rows"]
            ),
            "normalized_mychem_unresolved_examples": enrich_api_samples(
                api_payload["mychem_identity_results"]["normalized_inn"]["sample_unresolved_rows"]
            ),
        },
    }


def build_loader_design_payload(
    probe_payloads: dict[str, dict[str, Any]],
    validation_payload: dict[str, Any],
) -> dict[str, Any]:
    api_payload = probe_payloads["api_linkage"]
    vaccine_payload = probe_payloads["vaccine_identity"]
    metadata_payload = probe_payloads["metadata_and_devices"]
    validation_counts = {
        label: details["count"] for label, details in validation_payload["probes"].items()
    }
    return {
        "generated_at": iso_now(),
        "recommendation": {
            "winning_loader_strategy": "extend_existing_loader_for_apis_only",
            "separate_vaccine_loader_required": True,
            "device_catalog_in_scope": False,
            "evidence": [
                f"API normalized-INN MyChem phrase/exact resolution: {api_payload['mychem_identity_results']['normalized_inn']['phrase_or_exact_rate']}%",
                f"API component overlap with finished pharma: {api_payload['finished_pharma_overlap']['component_link_rate']}%",
                f"Best vaccine phrase/exact identity strategy: {vaccine_payload['winner']['phrase_or_exact_rate']}%",
                f"Vaccine immunization-schedule completeness: {metadata_payload['vaccine_specific_field_assessment']['immunization_schedule']['completeness']['percent']}%",
                f"Vaccine cold-chain completeness: {metadata_payload['vaccine_specific_field_assessment']['cold_chain_or_storage']['completeness']['percent']}%",
            ],
        },
        "product_type_filter_design": {
            "type": "who_product_type",
            "default_values": ["finished_pharma", "api"],
            "values": [
                {
                    "value": "finished_pharma",
                    "source_url": WHO_FINISHED_URL,
                    "include_by_default": True,
                },
                {
                    "value": "api",
                    "source_url": WHO_API_URL,
                    "include_by_default": True,
                },
                {
                    "value": "vaccine",
                    "source_url": WHO_VACCINES_URL,
                    "include_by_default": False,
                    "blocked_reason": "No stable INN bridge and missing vaccine-specific downstream metadata.",
                },
            ],
        },
        "proposed_who_prequalification_entry_extensions": {
            "status": "needed_for_api_support",
            "field_changes": [
                {
                    "field": "who_product_id",
                    "type": "Option<String>",
                    "applies_to": ["api"],
                    "reason": "API export keys rows with WHO Product ID instead of WHO reference number.",
                },
                {
                    "field": "grade",
                    "type": "Option<String>",
                    "applies_to": ["api"],
                    "reason": "API export includes Standard / other grade values not present in finished pharma.",
                },
                {
                    "field": "confirmation_document_date",
                    "type": "Option<String>",
                    "applies_to": ["api"],
                    "reason": "API export exposes a second dated regulatory milestone.",
                },
                {
                    "field": "dosage_form",
                    "type": "Option<String>",
                    "applies_to": ["finished_pharma", "api", "vaccine"],
                    "reason": "API and vaccine exports do not expose the finished-pharma dosage-form field.",
                },
                {
                    "field": "listing_basis",
                    "type": "Option<String>",
                    "applies_to": ["finished_pharma", "api", "vaccine"],
                    "reason": "API and vaccine exports do not carry finished-pharma listing-basis metadata.",
                },
                {
                    "field": "presentation",
                    "type": "Option<String>",
                    "applies_to": ["finished_pharma", "api"],
                    "reason": "API rows carry a canonical INN but not the finished-pharma presentation string.",
                },
            ],
        },
        "proposed_vaccine_surface": {
            "strategy": "separate_loader_or_schema_extension",
            "minimum_fields": [
                "vaccine_type",
                "commercial_name",
                "presentation",
                "dose_count",
                "manufacturer",
                "responsible_nra",
                "prequalification_date",
            ],
            "known_gaps": [
                "No WHO vaccine INN field",
                "No immunization schedule field",
                "No cold-chain or storage field",
            ],
            "validation_anchor_counts": validation_counts,
        },
        "device_catalog_decision": {
            "include_in_current_who_drug_pipeline": False,
            "source_url": WHO_DEVICES_URL,
            "reason": "Device catalog is a category-keyed product catalog with zero required-header overlap against the WHO drug contract.",
        },
    }


def build_contract_numbers(
    probe_payloads: dict[str, dict[str, Any]], stage_timings: dict[str, float]
) -> dict[str, Any]:
    schema_payload = probe_payloads["schema_comparison"]
    vaccine_payload = probe_payloads["vaccine_identity"]
    api_payload = probe_payloads["api_linkage"]
    metadata_payload = probe_payloads["metadata_and_devices"]
    return {
        "record_counts": schema_payload["record_counts"],
        "vaccine_identity": {
            "winning_strategy": vaccine_payload["winner"]["strategy"],
            "phrase_or_exact_rate": vaccine_payload["winner"]["phrase_or_exact_rate"],
            "exact_rate": vaccine_payload["winner"]["exact_rate"],
        },
        "api_identity": {
            "raw_inn_phrase_or_exact_rate": api_payload["mychem_identity_results"]["raw_inn"][
                "phrase_or_exact_rate"
            ],
            "normalized_inn_phrase_or_exact_rate": api_payload["mychem_identity_results"][
                "normalized_inn"
            ]["phrase_or_exact_rate"],
            "normalized_inn_exact_rate": api_payload["mychem_identity_results"]["normalized_inn"][
                "exact_match_rate"
            ],
            "finished_pharma_exact_overlap_rate": api_payload["finished_pharma_overlap"][
                "exact_inn_link_rate"
            ],
            "finished_pharma_component_overlap_rate": api_payload["finished_pharma_overlap"][
                "component_link_rate"
            ],
        },
        "vaccine_metadata": {
            "dose_count_percent": metadata_payload["vaccine_specific_field_assessment"]["dose_count"][
                "completeness"
            ]["percent"],
            "schedule_percent": metadata_payload["vaccine_specific_field_assessment"][
                "immunization_schedule"
            ]["completeness"]["percent"],
            "cold_chain_percent": metadata_payload["vaccine_specific_field_assessment"][
                "cold_chain_or_storage"
            ]["completeness"]["percent"],
        },
        "ticket_validation_counts": {
            key: value["count"] for key, value in metadata_payload["ticket_validation_samples"].items()
        },
        "build_timings_seconds": stage_timings,
    }


def diff_vaccine_identity_projection(
    baseline: dict[str, Any], current: dict[str, Any]
) -> list[dict[str, Any]]:
    diffs: list[dict[str, Any]] = []
    if current.get("row_count") != baseline.get("row_count"):
        diffs.append(
            {
                "path": "row_count",
                "baseline": baseline.get("row_count"),
                "current": current.get("row_count"),
            }
        )

    baseline_strategies = baseline.get("strategy_results", {})
    current_strategies = current.get("strategy_results", {})
    for strategy_name, baseline_metrics in baseline_strategies.items():
        current_metrics = current_strategies.get(strategy_name)
        if not isinstance(current_metrics, dict):
            diffs.append(
                {
                    "path": f"strategy_results.{strategy_name}",
                    "baseline": baseline_metrics,
                    "current": current_metrics,
                }
            )
            continue
        for metric_name, baseline_value in baseline_metrics.items():
            current_value = current_metrics.get(metric_name)
            path = f"strategy_results.{strategy_name}.{metric_name}"
            if current_value is None:
                diffs.append({"path": path, "baseline": baseline_value, "current": current_value})
                continue
            if metric_name == "rows":
                if current_value != baseline_value:
                    diffs.append(
                        {"path": path, "baseline": baseline_value, "current": current_value}
                    )
                continue
            if metric_name.startswith("no_"):
                if current_value > baseline_value:
                    diffs.append(
                        {"path": path, "baseline": baseline_value, "current": current_value}
                    )
                continue
            if current_value < baseline_value:
                diffs.append({"path": path, "baseline": baseline_value, "current": current_value})

    baseline_winner = baseline.get("winner", {})
    current_winner = current.get("winner", {})
    for metric_name in ("phrase_or_exact_rate", "exact_rate"):
        baseline_value = baseline_winner.get(metric_name)
        current_value = current_winner.get(metric_name)
        path = f"winner.{metric_name}"
        if baseline_value is None or current_value is None:
            diffs.append({"path": path, "baseline": baseline_value, "current": current_value})
            continue
        if current_value < baseline_value:
            diffs.append({"path": path, "baseline": baseline_value, "current": current_value})
    return diffs


def build_regression_control_payload() -> dict[str, Any]:
    ensure_work_dirs()
    baseline_payloads = {name: read_json(path) for name, path in PROBE_RESULT_PATHS.items()}
    baseline_projection = {
        name: build_probe_projection(name, payload) for name, payload in baseline_payloads.items()
    }
    current_payloads, stage_timings = run_probe_suite()
    current_projection = {
        name: build_probe_projection(name, payload) for name, payload in current_payloads.items()
    }

    comparisons: dict[str, Any] = {}
    all_passed = True
    for name in PROBE_RESULT_PATHS:
        comparison_mode = "exact_match"
        if name == "vaccine_identity":
            diffs = diff_vaccine_identity_projection(
                baseline_projection[name], current_projection[name]
            )
            comparison_mode = "match_or_beat"
        else:
            diffs = diff_payloads(
                baseline_projection[name], current_projection[name], drop_keys=frozenset()
            )
        passed = len(diffs) == 0
        all_passed = all_passed and passed
        comparisons[name] = {
            "passed": passed,
            "comparison_mode": comparison_mode,
            "baseline_projection_checksum": sha256_json(
                baseline_projection[name], drop_keys=frozenset()
            ),
            "current_projection_checksum": sha256_json(
                current_projection[name], drop_keys=frozenset()
            ),
            "mismatch_count": len(diffs),
            "diff_sample": diffs[:10],
            "baseline_projection": baseline_projection[name],
            "current_projection": current_projection[name],
        }

    payload = {
        "generated_at": iso_now(),
        "rule": {
            "type": "probe_specific_match_or_beat",
            "note": "Regression control keeps exact metric projection matching for schema, API, and metadata probes. Vaccine identity is evaluated as match-or-beat against the explore projection so optimize-step improvements can pass without relaxing the other guards.",
        },
        "baseline_paths": {name: str(path.resolve()) for name, path in PROBE_RESULT_PATHS.items()},
        "stage_timings_seconds": stage_timings,
        "source_snapshots": {
            "finished_pharma": source_snapshot(WHO_FINISHED_URL),
            "vaccines": source_snapshot(WHO_VACCINES_URL),
            "apis": source_snapshot(WHO_API_URL),
            "devices": source_snapshot(WHO_DEVICES_URL),
        },
        "comparisons": comparisons,
        "overall_passed": all_passed,
    }
    write_json(REGRESSION_CONTROL_PATH, payload)
    if not all_passed:
        raise RuntimeError("WHO exploit regression control failed")
    return payload


def write_regression_control_result() -> dict[str, Any]:
    return build_regression_control_payload()


def build_full_scale_results() -> WhoFullScaleArtifacts:
    ensure_work_dirs()
    timer_started = time.perf_counter()
    probe_payloads, probe_timings = run_probe_suite()

    validation_payload, validation_seconds = time_call(build_validation_payload)
    write_json(VALIDATION_PATH, validation_payload)

    sample_payload, sample_seconds = time_call(
        build_sample_records_payload, probe_payloads, validation_payload
    )
    write_json(SAMPLES_PATH, sample_payload)

    design_payload, design_seconds = time_call(
        build_loader_design_payload, probe_payloads, validation_payload
    )
    write_json(DESIGN_PATH, design_payload)

    total_seconds = round_seconds(time.perf_counter() - timer_started)
    stage_timings = {
        **probe_timings,
        "validation_seconds": validation_seconds,
        "sample_records_seconds": sample_seconds,
        "design_payload_seconds": design_seconds,
        "total_seconds": total_seconds,
    }
    contract_numbers = build_contract_numbers(probe_payloads, stage_timings)
    artifact_paths = {
        "probe_summary": str(PROBE_SUMMARY_PATH.resolve()),
        "regression_control": str(REGRESSION_CONTROL_PATH.resolve()),
        "validation": str(VALIDATION_PATH.resolve()),
        "sample_records": str(SAMPLES_PATH.resolve()),
        "loader_design": str(DESIGN_PATH.resolve()),
        "full_scale_results": str(FULL_SCALE_RESULTS_PATH.resolve()),
    }
    full_scale_payload = {
        "generated_at": iso_now(),
        "full_scale_definition": "All three WHO lists plus the immunization-device catalog loaded, cross-referenced, validated, and summarized with persistent artifacts.",
        "cache_state": current_cache_state(),
        "source_snapshots": {
            "finished_pharma": source_snapshot(WHO_FINISHED_URL),
            "vaccines": source_snapshot(WHO_VACCINES_URL),
            "apis": source_snapshot(WHO_API_URL),
            "devices": source_snapshot(WHO_DEVICES_URL),
        },
        "contract_numbers": contract_numbers,
        "recommendation": design_payload["recommendation"],
        "artifact_paths": artifact_paths,
    }
    write_json(FULL_SCALE_RESULTS_PATH, full_scale_payload)
    return WhoFullScaleArtifacts(
        probe_payloads=probe_payloads,
        validation_payload=validation_payload,
        sample_records_payload=sample_payload,
        loader_design_payload=design_payload,
        full_scale_payload=full_scale_payload,
        stage_timings=stage_timings,
        contract_numbers=contract_numbers,
        artifact_paths=artifact_paths,
    )
