from __future__ import annotations

import argparse
import json
import re
import sys
from collections import Counter
from datetime import datetime, timezone
from pathlib import Path
from zoneinfo import ZoneInfo


ROOT = Path(__file__).resolve().parents[1]
WORKSPACE = ROOT.parent
DEFAULT_DATA_DIR = WORKSPACE / "medical-tools-extension-data" / "consultations"
DEFAULT_TIMEZONE = "Europe/Kiev"
SERVICE_TEXT_PATTERNS = [
    re.compile(r"дослівн[а-яіїєґ'\s-]*транскрипц", re.I),
    re.compile(r"контекст[:\s-]*консультац", re.I),
    re.compile(r"зберігай\s+медичн", re.I),
    re.compile(r"не\s+додавай\s+дан", re.I),
    re.compile(r"не\s+додай\s+нічого", re.I),
    re.compile(r"об['ʼ]єктивн[а-яіїєґ'\s,-]*статус\s+і\s+рекомендац", re.I),
]
NOISE_TEXT_PATTERNS = [
    re.compile(r"^дякую[.!?…]*$", re.I),
    re.compile(r"^дякую\s+за\s+(?:перегляд|увагу)[.!?…]*$", re.I),
    re.compile(r"^спасибі[.!?…]*$", re.I),
    re.compile(r"^дитячий\s+невролог[.!?…]*$", re.I),
    re.compile(r"^консультац[а-яіїєґ']*\s+дитяч[а-яіїєґ']*\s+невролог[а-яіїєґ']*[.!?…]*$", re.I),
    re.compile(r"^(?:звучить\s+музика|лунає\s+музика|музика|сміх)[.!?…]*$", re.I),
]


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Audit the latest Parallax Health consultation session without printing PHI."
    )
    parser.add_argument("--data-dir", default=str(DEFAULT_DATA_DIR), help="Consultations data directory.")
    parser.add_argument("--session", default="", help="Specific session id. Defaults to newest session by file mtime.")
    parser.add_argument(
        "--source",
        default="",
        help="Optional session source filter, for example microphone, mic, file_test, readiness_microphone_test.",
    )
    parser.add_argument(
        "--around-local-time",
        default="",
        metavar="HH:MM",
        help="Pick the session whose finish/start time is closest to this local time.",
    )
    parser.add_argument(
        "--date",
        default="",
        metavar="YYYY-MM-DD",
        help="Local date for --around-local-time. Defaults to today in --timezone.",
    )
    parser.add_argument("--timezone", default=DEFAULT_TIMEZONE, help="Timezone for local time matching.")
    parser.add_argument("--strict", action="store_true", help="Exit non-zero when the audit verdict is FAIL.")
    args = parser.parse_args()

    data_dir = Path(args.data_dir).expanduser()
    session_dir = select_session_dir(data_dir, args)
    if not session_dir:
        raise SystemExit(f"No consultation sessions found in {data_dir}")

    report = build_report(session_dir)
    print(format_report(report))
    return 1 if args.strict and report["verdict"] == "FAIL" else 0


def select_session_dir(data_dir: Path, args: argparse.Namespace) -> Path | None:
    source_filter = normalize_source_filter(getattr(args, "source", ""))
    if args.session:
        session_dir = data_dir / args.session
        if source_filter and not session_source_matches(session_dir, source_filter):
            return None
        return session_dir
    if args.around_local_time:
        return session_dir_around_local_time(
            data_dir,
            args.around_local_time,
            args.date,
            args.timezone,
            source_filter,
        )
    return latest_session_dir(data_dir, source_filter)


def latest_session_dir(data_dir: Path, source_filter: str = "") -> Path | None:
    if not data_dir.is_dir():
        return None
    normalized_source = normalize_source_filter(source_filter)
    candidates = [
        path
        for path in data_dir.iterdir()
        if path.is_dir() and session_source_matches(path, normalized_source)
    ]
    if not candidates:
        return None
    return max(candidates, key=latest_mtime)


def session_dir_around_local_time(
    data_dir: Path,
    time_text: str,
    date_text: str,
    timezone_name: str,
    source_filter: str = "",
) -> Path | None:
    if not data_dir.is_dir():
        return None
    target_time = parse_hh_mm(time_text)
    tz = ZoneInfo(timezone_name)
    local_date = parse_local_date(date_text, tz)
    target_dt = datetime.combine(local_date, target_time, tzinfo=tz)
    normalized_source = normalize_source_filter(source_filter)
    candidates: list[tuple[float, Path]] = []
    for session_dir in data_dir.iterdir():
        if not session_dir.is_dir():
            continue
        if not session_source_matches(session_dir, normalized_source):
            continue
        best_time = best_session_local_time(session_dir, tz)
        if not best_time or best_time.date() != local_date:
            continue
        candidates.append((abs((best_time - target_dt).total_seconds()), session_dir))
    if not candidates:
        return None
    return min(candidates, key=lambda item: (item[0], latest_mtime(item[1])))[1]


def normalize_source_filter(value: object) -> str:
    source = str(value or "").strip().lower()
    aliases = {
        "mic": "microphone",
        "live": "microphone",
        "recording": "microphone",
        "ready": "readiness_microphone_test",
        "readiness": "readiness_microphone_test",
        "file": "file_test",
    }
    return aliases.get(source, source)


def session_source_matches(session_dir: Path, source_filter: str) -> bool:
    if not source_filter:
        return True
    source = normalize_source_filter(session_source(session_dir))
    if source_filter == "file_test":
        return source.startswith("file_test")
    return source == source_filter


def session_source(session_dir: Path) -> str:
    session = load_json(session_dir / "session.json")
    status = load_json(session_dir / "status.json")
    metrics = load_json(session_dir / "metrics.json")
    return first_text(
        status.get("sessionSource"),
        metrics.get("sessionSource"),
        status.get("source"),
        metrics.get("source"),
        session.get("source"),
        "microphone",
    )


def parse_hh_mm(value: str):
    try:
        return datetime.strptime(value.strip(), "%H:%M").time()
    except ValueError as error:
        raise SystemExit("--around-local-time must use HH:MM, for example 08:52.") from error


def parse_local_date(value: str, tz: ZoneInfo):
    if value:
        try:
            return datetime.strptime(value.strip(), "%Y-%m-%d").date()
        except ValueError as error:
            raise SystemExit("--date must use YYYY-MM-DD, for example 2026-06-10.") from error
    return datetime.now(tz).date()


def best_session_local_time(session_dir: Path, tz: ZoneInfo) -> datetime | None:
    session = load_json(session_dir / "session.json")
    status = load_json(session_dir / "status.json")
    timestamps = [
        status.get("clientFinishedAt"),
        session.get("clientFinishedAt"),
        status.get("finishedAt"),
        session.get("finishedAt"),
        status.get("clientStartedAt"),
        session.get("clientStartedAt"),
        status.get("startedAt"),
        session.get("startedAt"),
    ]
    parsed = [parse_iso_datetime(value, tz) for value in timestamps if value]
    parsed = [value for value in parsed if value is not None]
    return parsed[0] if parsed else None


def parse_iso_datetime(value: object, tz: ZoneInfo) -> datetime | None:
    text = str(value or "").strip()
    if not text:
        return None
    try:
        parsed = datetime.fromisoformat(text.replace("Z", "+00:00"))
    except ValueError:
        return None
    if parsed.tzinfo is None:
        parsed = parsed.replace(tzinfo=timezone.utc)
    return parsed.astimezone(tz)


def latest_mtime(path: Path) -> float:
    newest = path.stat().st_mtime
    for child in path.rglob("*"):
        try:
            if child.is_file():
                newest = max(newest, child.stat().st_mtime)
        except OSError:
            continue
    return newest


def build_report(session_dir: Path) -> dict:
    session = load_json(session_dir / "session.json")
    status = load_json(session_dir / "status.json")
    metrics = load_json(session_dir / "metrics.json")
    transcript_path = session_dir / "transcript.txt"
    speaker_path = session_dir / "transcript-speakers.txt"
    draft_path = session_dir / "lakmus-draft.json"
    chunks = session.get("chunks") if isinstance(session.get("chunks"), list) else []
    errors = session.get("errors") if isinstance(session.get("errors"), list) else []
    segments = session.get("segments") if isinstance(session.get("segments"), list) else []

    chunk_statuses = Counter(str(chunk.get("status") or "") for chunk in chunks if isinstance(chunk, dict))
    mime_types = Counter(str(chunk.get("mimeType") or "") for chunk in chunks if isinstance(chunk, dict))
    error_messages = [str(item.get("message") or item.get("error") or "") for item in errors if isinstance(item, dict)]
    if not error_messages:
        error_messages = [str(chunk.get("error") or "") for chunk in chunks if isinstance(chunk, dict) and chunk.get("error")]

    transcript_bytes = safe_size(transcript_path)
    speaker_bytes = safe_size(speaker_path)
    draft = load_json(draft_path)
    draft_fields = draft.get("fields") if isinstance(draft.get("fields"), list) else []
    draft_fields_with_text = sum(1 for field in draft_fields if str(field.get("value") or "").strip())
    client_extension_version = first_text(
        status.get("clientExtensionVersion"),
        metrics.get("clientExtensionVersion"),
        session.get("clientExtensionVersion"),
        session.get("extensionVersion"),
    )
    client_audio_pipeline = first_text(
        status.get("clientAudioPipeline"),
        metrics.get("clientAudioPipeline"),
        session.get("clientAudioPipeline"),
    )
    client_live_chunk_policy = first_text(
        status.get("clientLiveChunkPolicy"),
        metrics.get("clientLiveChunkPolicy"),
        session.get("clientLiveChunkPolicy"),
    )
    session_source = first_text(
        status.get("sessionSource"),
        metrics.get("sessionSource"),
        status.get("source"),
        metrics.get("source"),
        session.get("source"),
    )
    is_file_test = session_source.startswith("file_test")

    findings: list[str] = []
    verdict = "PASS"
    technical_verdict = "PASS"
    clinical_verdict = "PASS"
    clinical_state = "not_evaluated"
    error_count = int(status.get("errorCount") or metrics.get("errorChunkCount") or chunk_statuses.get("error", 0) or 0)
    chunk_count = int(status.get("chunkCount") or metrics.get("chunkCount") or len(chunks) or 0)
    done_count = int(status.get("doneChunkCount") or metrics.get("doneChunkCount") or chunk_statuses.get("done", 0) or 0)
    fallback_whisper_counts = whisper_discard_counts(session_dir)
    raw_whisper_segment_count = int(metrics.get("rawWhisperSegmentCount") or sum_numeric(chunks, "rawWhisperSegmentCount") or fallback_whisper_counts["raw"] or 0)
    discarded_whisper_segment_count = int(metrics.get("discardedWhisperSegmentCount") or sum_numeric(chunks, "discardedWhisperSegmentCount") or fallback_whisper_counts["discarded"] or 0)
    useful_whisper_segment_count = int(metrics.get("usefulWhisperSegmentCount") or sum_numeric(chunks, "usefulWhisperSegmentCount") or fallback_whisper_counts["useful"] or 0)

    if chunk_count == 0:
        verdict = "WARN"
        technical_verdict = max_verdict(technical_verdict, "WARN")
        findings.append("No chunks recorded yet.")
    if error_count:
        verdict = "FAIL"
        technical_verdict = "FAIL"
        findings.append(f"{error_count} chunk errors detected.")
    if any("webm" in item.lower() for item in mime_types) and not is_file_test:
        if error_count:
            verdict = "FAIL"
            technical_verdict = "FAIL"
        else:
            verdict = max_verdict(verdict, "WARN")
            technical_verdict = max_verdict(technical_verdict, "WARN")
        findings.append("WebM chunks are present; current live recorder should normally produce audio/wav.")
    if chunk_count and not client_live_chunk_policy and not is_file_test:
        technical_verdict = max_verdict(technical_verdict, "WARN")
        verdict = max_verdict(verdict, "WARN")
        findings.append("Client live chunk policy is missing; this is likely an older recorder build.")
    elif client_live_chunk_policy and client_live_chunk_policy not in {"wav_only", "not_live_file_test"}:
        verdict = max_verdict(verdict, "WARN")
        technical_verdict = max_verdict(technical_verdict, "WARN")
        findings.append(f"Client live chunk policy is {client_live_chunk_policy}; expected wav_only for current recorder.")
    if any("ebml" in message.lower() for message in error_messages):
        verdict = "FAIL"
        technical_verdict = "FAIL"
        findings.append("EBML header parsing failure detected; this matches the old sliced-WebM bug.")
    if chunk_count and done_count == 0:
        verdict = "FAIL"
        technical_verdict = "FAIL"
        findings.append("No chunks completed transcription.")
    if done_count and transcript_bytes == 0:
        verdict = max_verdict(verdict, "WARN")
        technical_verdict = max_verdict(technical_verdict, "WARN")
        if discarded_whisper_segment_count:
            if useful_whisper_segment_count:
                findings.append("Chunks completed; some Whisper fragments were filtered before transcript output, and no usable transcript survived normalization.")
            else:
                findings.append("Chunks completed; Whisper returned only short noise/service fragments that were filtered before transcript output.")
        else:
            findings.append("Chunks completed but transcript is empty; likely no speech or weak audio.")
    draft_eligible_count = int(metrics.get("draftEligibleSegmentCount") or 0)
    clinical_signal_count = int(metrics.get("clinicalSignalSegmentCount") or 0)
    draft_empty_reason = clean_text(metrics.get("draftEmptyReason"))
    educational_content = bool(metrics.get("educationalContent")) or draft_empty_reason == "educational_content"
    if transcript_bytes and draft_fields_with_text == 0:
        if educational_content:
            clinical_verdict = "NONE"
            clinical_state = "educational_content_no_lakmus_fields"
            findings.append("Transcript exists, but the session looks like educational/test content; Lakmus fields should stay empty.")
        elif draft_eligible_count == 0:
            clinical_verdict = "NONE"
            clinical_state = "no_clinical_field_text"
            findings.append("Transcript exists, but no clinical draft-eligible segments were detected; Lakmus fields should stay empty.")
        else:
            clinical_verdict = "WARN"
            clinical_state = "clinical_segments_without_fields"
            verdict = max_verdict(verdict, "WARN")
            findings.append("Clinical draft-eligible segments exist, but Lakmus fields are empty; review mapper output.")
    elif draft_fields_with_text:
        clinical_verdict = "PASS"
        clinical_state = "fields_ready"
    elif done_count:
        clinical_verdict = "NONE"
        clinical_state = "no_transcript_text"
    if not findings:
        findings.append("No technical blocker detected in saved session artifacts.")
    operator_action = operator_action_for(technical_verdict, clinical_verdict, clinical_state)

    return {
        "verdict": verdict,
        "technicalVerdict": technical_verdict,
        "clinicalVerdict": clinical_verdict,
        "clinicalState": clinical_state,
        "operatorAction": operator_action,
        "sessionId": session_dir.name,
        "sessionDir": str(session_dir),
        "checkedAt": datetime.now(timezone.utc).replace(microsecond=0).isoformat(),
        "status": status.get("status") or session.get("status") or "",
        "processingState": status.get("processingState") or session.get("processingState") or "",
        "sessionSource": session_source,
        "clientExtensionVersion": client_extension_version,
        "clientAudioPipeline": client_audio_pipeline,
        "clientLiveChunkPolicy": client_live_chunk_policy,
        "startedAt": status.get("startedAt") or session.get("startedAt") or "",
        "finishedAt": status.get("finishedAt") or session.get("finishedAt") or "",
        "chunkCount": chunk_count,
        "doneChunkCount": done_count,
        "errorChunkCount": error_count,
        "pendingChunkCount": int(status.get("pendingChunkCount") or metrics.get("pendingChunkCount") or 0),
        "rawWhisperSegmentCount": raw_whisper_segment_count,
        "usefulWhisperSegmentCount": useful_whisper_segment_count,
        "discardedWhisperSegmentCount": discarded_whisper_segment_count,
        "segmentCount": int(status.get("segmentCount") or metrics.get("segmentCount") or len(segments) or 0),
        "clinicalSignalSegmentCount": clinical_signal_count,
        "draftEligibleSegmentCount": draft_eligible_count,
        "draftEmptyReason": draft_empty_reason,
        "educationalContent": educational_content,
        "transcriptBytes": transcript_bytes,
        "speakerTranscriptBytes": speaker_bytes,
        "draftFieldsWithText": draft_fields_with_text,
        "mimeTypes": dict(sorted(mime_types.items())),
        "chunkStatuses": dict(sorted(chunk_statuses.items())),
        "avgInputRmsDb": metrics.get("avgInputRmsDb", ""),
        "avgActiveFrameRatio": metrics.get("avgActiveFrameRatio", ""),
        "firstTranscriptLatencySeconds": metrics.get("firstTranscriptLatencySeconds", ""),
        "sampleError": sanitize_error(error_messages[0]) if error_messages else "",
        "findings": findings,
    }


def max_verdict(current: str, candidate: str) -> str:
    order = {"PASS": 0, "WARN": 1, "FAIL": 2}
    return candidate if order[candidate] > order[current] else current


def first_text(*values: object) -> str:
    for value in values:
        text = str(value or "").strip()
        if text:
            return text
    return ""


def sum_numeric(items: list, key: str) -> int:
    total = 0
    for item in items:
        if not isinstance(item, dict):
            continue
        try:
            total += int(float(item.get(key) or 0))
        except (TypeError, ValueError):
            continue
    return total


def whisper_discard_counts(session_dir: Path) -> dict:
    counts = {"raw": 0, "useful": 0, "discarded": 0}
    whisper_dir = session_dir / "whisper"
    if not whisper_dir.is_dir():
        return counts
    for path in sorted(whisper_dir.glob("*.json")):
        data = load_json(path)
        texts = whisper_texts(data)
        for text in texts:
            counts["raw"] += 1
            if is_discardable_text(text):
                counts["discarded"] += 1
            else:
                counts["useful"] += 1
    return counts


def whisper_texts(data: dict) -> list[str]:
    segments = data.get("segments") if isinstance(data.get("segments"), list) else []
    if segments:
        return [clean_text(segment.get("text") if isinstance(segment, dict) else "") for segment in segments if clean_text(segment.get("text") if isinstance(segment, dict) else "")]
    text = clean_text(data.get("text"))
    return [text] if text else []


def clean_text(value: object) -> str:
    return " ".join(str(value or "").split()).strip()


def is_discardable_text(value: object) -> bool:
    text = clean_text(value).lower()
    if not text:
        return False
    if any(pattern.search(text) for pattern in SERVICE_TEXT_PATTERNS):
        return True
    words = text.split()
    return len(words) <= 5 and any(pattern.search(text) for pattern in NOISE_TEXT_PATTERNS)


def operator_action_for(technical_verdict: str, clinical_verdict: str, clinical_state: str) -> str:
    if technical_verdict == "FAIL":
        return "technical_failure_do_not_rely_on_recording"
    if technical_verdict == "WARN":
        return "review_technical_warnings_before_clinical_use"
    if clinical_verdict == "WARN":
        return "review_mapper_fields_before_lakmus_fill"
    if clinical_state == "educational_content_no_lakmus_fields":
        return "educational_or_test_audio_transcribed_no_lakmus_fields"
    if clinical_state == "no_clinical_field_text":
        return "recording_transcribed_no_lakmus_fields_expected"
    if clinical_state == "fields_ready":
        return "review_fields_then_fill_visible_lakmus"
    return "recording_technical_path_ok"


def format_report(report: dict) -> str:
    lines = [
        "Parallax Health consultation QA audit",
        "privacy: no audio, no transcript text, no Lakmus field text, no license key",
        f"checkedAt: {report['checkedAt']}",
        f"verdict: {report['verdict']}",
        f"technicalVerdict: {report['technicalVerdict']}",
        f"clinicalVerdict: {report['clinicalVerdict']}",
        f"clinicalState: {report['clinicalState']}",
        f"operatorAction: {report['operatorAction']}",
        f"sessionId: {report['sessionId']}",
        f"status: {report['status']}",
        f"processingState: {report['processingState']}",
        f"sessionSource: {report['sessionSource']}",
        f"clientExtensionVersion: {report['clientExtensionVersion']}",
        f"clientAudioPipeline: {report['clientAudioPipeline']}",
        f"clientLiveChunkPolicy: {report['clientLiveChunkPolicy']}",
        f"startedAt: {report['startedAt']}",
        f"finishedAt: {report['finishedAt']}",
        f"chunkCount: {report['chunkCount']}",
        f"doneChunkCount: {report['doneChunkCount']}",
        f"errorChunkCount: {report['errorChunkCount']}",
        f"pendingChunkCount: {report['pendingChunkCount']}",
        f"rawWhisperSegmentCount: {report['rawWhisperSegmentCount']}",
        f"usefulWhisperSegmentCount: {report['usefulWhisperSegmentCount']}",
        f"discardedWhisperSegmentCount: {report['discardedWhisperSegmentCount']}",
        f"segmentCount: {report['segmentCount']}",
        f"clinicalSignalSegmentCount: {report['clinicalSignalSegmentCount']}",
        f"draftEligibleSegmentCount: {report['draftEligibleSegmentCount']}",
        f"transcriptBytes: {report['transcriptBytes']}",
        f"speakerTranscriptBytes: {report['speakerTranscriptBytes']}",
        f"draftFieldsWithText: {report['draftFieldsWithText']}",
        f"mimeTypes: {json.dumps(report['mimeTypes'], ensure_ascii=False)}",
        f"chunkStatuses: {json.dumps(report['chunkStatuses'], ensure_ascii=False)}",
        f"avgInputRmsDb: {report['avgInputRmsDb']}",
        f"avgActiveFrameRatio: {report['avgActiveFrameRatio']}",
        f"firstTranscriptLatencySeconds: {report['firstTranscriptLatencySeconds']}",
    ]
    if report["sampleError"]:
        lines.append(f"sampleError: {report['sampleError']}")
    lines.append("findings:")
    lines.extend(f"- {item}" for item in report["findings"])
    return "\n".join(lines).rstrip() + "\n"


def load_json(path: Path) -> dict:
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except Exception:
        return {}


def safe_size(path: Path) -> int:
    try:
        return path.stat().st_size if path.is_file() else 0
    except OSError:
        return 0


def sanitize_error(value: str) -> str:
    value = " ".join(str(value or "").split())
    replacements = [
        (str(WORKSPACE), "<workspace>"),
        (str(Path.home()), "<home>"),
    ]
    for needle, replacement in replacements:
        value = value.replace(needle, replacement)
    return value[:320]


if __name__ == "__main__":
    sys.exit(main())
