agents.py

from __future__ import annotations

import json
import re
from pathlib import Path
from typing import Any, Dict

from loguru import logger

from config import LLM
from schemas import ClaimsState, SchedulingState
from tools import (
    appointment_option_ranker,
    attachment_manifest_generator,
    authorization_lookup,
    canonical_claim_schema_mapper,
    claim_field_extractor,
    claim_packet_uploader,
    denial_risk_classifier,
    duplicate_claim_checker,
    edi_like_json_parser,
    exception_similarity_rag_retriever,
    human_review_routing_tool,
    member_benefit_lookup,
    mock_eligibility_lookup,
    mock_provider_npi_registry_lookup,
    policy_benefit_rag_retriever,
    provider_availability_lookup,
    provider_note_parser,
    provider_specialty_rag_retriever,
    referral_lookup,
    required_field_validator,
    schedule_readiness_checker,
    scheduling_request_parser,
    scheduling_summary_writer,
    specialist_location_lookup,
)


def _merge_messages(state: Dict[str, Any], message: str) -> list[str]:
    return state.get("messages", []) + [message]


def _extract_json_object(text: str) -> Dict[str, Any]:
    """Best-effort parser for JSON-only LLM responses."""
    try:
        return json.loads(text)
    except Exception:
        pass

    match = re.search(r"\{.*\}", text, flags=re.DOTALL)
    if match:
        try:
            return json.loads(match.group(0))
        except Exception:
            logger.warning(f"Could not parse LLM JSON payload: {text[:300]}")
    return {}


def _llm_json_decision(prompt: str, fallback: Dict[str, Any]) -> Dict[str, Any]:
    """Call the shared LLM when enabled; otherwise return deterministic fallback."""
    if LLM is None:
        logger.debug(f"LLM disabled; using fallback decision: {fallback}")
        return fallback
    try:
        logger.info("Calling LLM for JSON decision")
        response = LLM.invoke(prompt)
        content = getattr(response, "content", str(response))
        parsed = _extract_json_object(content)
        return parsed or fallback
    except Exception as exc:
        logger.warning(f"LLM decision failed; using fallback. Error: {exc}")
        return fallback


# ---------------- Claims graph nodes ----------------


def claims_intake_node(state: ClaimsState) -> ClaimsState:
    packet_path = state["selected_packet_path"]
    logger.info(f"Claims intake node packet={packet_path}")
    claim_packet_uploader.invoke({"packet_path": packet_path})
    manifest = attachment_manifest_generator.invoke({"packet_path": packet_path})
    return {
        **state,
        "case_id": Path(packet_path).name,
        "packet_manifest": manifest,
        "messages": _merge_messages(state, "Intake complete"),
    }


def claims_extraction_node(state: ClaimsState) -> ClaimsState:
    logger.info("Claims extraction node")
    manifest = state.get("packet_manifest", {}).get("attachments", [])
    extracted_parts = []
    for item in manifest:
        path = item["file_path"]
        if item["file_type"] == "json":
            parsed = edi_like_json_parser.invoke({"json_path": path})
            extracted_parts.append(parsed.get("claim", {}))
        elif item["file_type"] == "text":
            note = provider_note_parser.invoke({"note_text_or_path": path})
            extracted_parts.append(note)

    combined_text = json.dumps(extracted_parts)
    extracted = claim_field_extractor.invoke({"text_or_json": combined_text})

    # Optional LLM extraction validation: useful for ambiguous note/OCR content.
    extraction_review = _llm_json_decision(
        prompt=f"""
You are validating extracted healthcare claim fields. Return JSON only.

Return this shape:
{{
  "extraction_confidence": "high|medium|low",
  "missing_or_ambiguous_fields": [],
  "reason": "short explanation"
}}

Extracted payload:
{json.dumps(extracted, indent=2)}
""",
        fallback={
            "extraction_confidence": (
                "high" if extracted.get("extracted", {}).get("claim_id") else "medium"
            ),
            "missing_or_ambiguous_fields": [],
            "reason": "Deterministic extraction review; LLM disabled or unavailable.",
        },
    )

    retry_count = int(state.get("extraction_retry_count", 0)) + 1
    return {
        **state,
        "extracted": extracted,
        "extraction_review": extraction_review,
        "extraction_retry_count": retry_count,
        "messages": _merge_messages(state, "Extraction complete"),
    }


def claims_normalization_node(state: ClaimsState) -> ClaimsState:
    logger.info("Claims normalization node")
    canonical = canonical_claim_schema_mapper.invoke(
        {"extracted_payload_json": json.dumps(state["extracted"])}
    )
    return {
        **state,
        "canonical_claim": canonical["canonical_claim"],
        "messages": _merge_messages(state, "Normalization complete"),
    }


def claims_validation_node(state: ClaimsState) -> ClaimsState:
    logger.info("Claims validation node")
    claim = state["canonical_claim"]
    required = required_field_validator.invoke(
        {"canonical_claim_json": json.dumps(claim)}
    )
    elig = mock_eligibility_lookup.invoke(
        {
            "member_id": claim.get("member", {}).get("member_id", ""),
            "service_date": (claim.get("service", {}).get("dates") or [""])[0],
            "claim_json": json.dumps(claim),
        }
    )
    provider = mock_provider_npi_registry_lookup.invoke(
        {
            "npi": claim.get("provider", {}).get("npi", ""),
            "claim_json": json.dumps(claim),
        }
    )
    duplicate = duplicate_claim_checker.invoke(
        {"canonical_claim_json": json.dumps(claim)}
    )
    validation = {
        **required,
        "eligible": elig.get("eligible"),
        "provider_valid": provider.get("valid"),
        "network_status": provider.get("network_status"),
        "in_network": provider.get("in_network"),
        "duplicate_risk": duplicate.get("duplicate_risk"),
        "duplicate_partner": duplicate.get("duplicate_partner"),
        "eligibility": elig,
        "provider": provider,
    }
    return {
        **state,
        "validation_results": validation,
        "messages": _merge_messages(state, "Validation complete"),
    }


def claims_validation_router(state: ClaimsState) -> str:
    """Conditional edge after validation."""
    validation = state.get("validation_results", {})
    retry_count = int(state.get("extraction_retry_count", 0))

    if validation.get("missing_fields"):
        if retry_count < 2:
            logger.info(
                f"Claims router: missing fields {validation.get('missing_fields')}; "
                f"retry extraction (attempt {retry_count}/1)"
            )
            return "retry_extraction"
        logger.warning(
            f"Claims router: missing fields {validation.get('missing_fields')} "
            f"after {retry_count} extraction attempts; continuing without retry"
        )

    if validation.get("eligible") is False or validation.get("provider_valid") is False:
        logger.info("Claims router: hard validation failure; route to exception")
        return "exception"

    logger.info("Claims router: continue to policy/similarity retrieval")
    return "policy_similarity"


def claims_policy_similarity_node(state: ClaimsState) -> ClaimsState:
    logger.info("Claims policy/similarity node")
    query = json.dumps(
        {
            "claim": state.get("canonical_claim"),
            "validation": state.get("validation_results"),
        }
    )

    # LLM decides which retrieval tools are necessary. Fallback stays safe and predictable.
    validation = state.get("validation_results", {})
    fallback_decision = {
        "use_policy_rag": True,
        "use_exception_similarity_rag": bool(
            validation.get("duplicate_risk")
            or validation.get("missing_fields")
            or validation.get("eligible") is False
        ),
        "reason": "Fallback rule: always check policy; use exception similarity when validation risk is present.",
    }
    rag_decision = _llm_json_decision(
        prompt=f"""
Decide which retrieval tools are needed for this claim. Return JSON only.

Allowed retrieval tools:
- policy_benefit_rag_retriever: payer rules, coding guidance, medical-necessity criteria, SOPs
- exception_similarity_rag_retriever: similar prior resolved exceptions

Return this shape:
{{
  "use_policy_rag": true,
  "use_exception_similarity_rag": false,
  "reason": "short explanation"
}}

Claim:
{json.dumps(state.get('canonical_claim'), indent=2)}

Validation:
{json.dumps(state.get('validation_results'), indent=2)}
""",
        fallback=fallback_decision,
    )

    policy = {
        "ok": True,
        "skipped": True,
        "results": [],
        "reason": rag_decision.get("reason"),
    }
    exceptions = {
        "ok": True,
        "skipped": True,
        "results": [],
        "reason": rag_decision.get("reason"),
    }

    if rag_decision.get("use_policy_rag"):
        policy = policy_benefit_rag_retriever.invoke({"query": query, "k": 5})
    if rag_decision.get("use_exception_similarity_rag"):
        exceptions = exception_similarity_rag_retriever.invoke({"query": query, "k": 5})

    policy_analysis = _llm_json_decision(
        prompt=f"""
Analyze the retrieved claim policy and exception context. Return JSON only.

Return this shape:
{{
  "likely_denial_risks": [],
  "recommended_next_steps": [],
  "explanation": "short audit-friendly summary"
}}

Claim:
{json.dumps(state.get('canonical_claim'), indent=2)}

Policy results:
{json.dumps(policy, indent=2)[:6000]}

Exception results:
{json.dumps(exceptions, indent=2)[:6000]}
""",
        fallback={
            "likely_denial_risks": [],
            "recommended_next_steps": [],
            "explanation": "LLM disabled or unavailable; deterministic tools supplied retrieval results only.",
        },
    )

    return {
        **state,
        "rag_decision": rag_decision,
        "policy_results": policy,
        "exception_results": exceptions,
        "policy_analysis": policy_analysis,
        "messages": _merge_messages(state, "Policy and similarity retrieval complete"),
    }


def claims_exception_node(state: ClaimsState) -> ClaimsState:
    logger.info("Claims exception node")
    combined_rag = {
        "policy": state.get("policy_results"),
        "exceptions": state.get("exception_results"),
        "policy_analysis": state.get("policy_analysis"),
    }
    risk = denial_risk_classifier.invoke(
        {
            "validation_results_json": json.dumps(state.get("validation_results", {})),
            "rag_results_json": json.dumps(combined_rag),
        }
    )
    route = human_review_routing_tool.invoke({"denial_risk_json": json.dumps(risk)})

    llm_decision = _llm_json_decision(
        prompt=f"""
You are the exception decision agent for a payer claim. Return JSON only.

Return this shape:
{{
  "risk_level": "low|medium|high",
  "recommended_route": "clean_pass_auto_normalization|claims_ops_exception_review|claims_ops_duplicate_review|prior_auth_exception_review",
  "explanation": "brief audit-friendly explanation",
  "confidence": 0.0
}}

Claim:
{json.dumps(state.get('canonical_claim'), indent=2)}

Validation:
{json.dumps(state.get('validation_results'), indent=2)}

Deterministic risk:
{json.dumps(risk, indent=2)}

Deterministic route:
{json.dumps(route, indent=2)}

Retrieved context summary:
{json.dumps(combined_rag, indent=2)[:6000]}
""",
        fallback={
            "risk_level": risk.get("risk_level", "low"),
            "recommended_route": route.get("route"),
            "explanation": "Deterministic exception routing; LLM disabled or unavailable.",
            "confidence": 0.75,
        },
    )

    summary = {
        "claim_id": state.get("canonical_claim", {}).get("claim_id"),
        "risk": risk,
        "route": route,
        "llm_decision": llm_decision,
        "rag_decision": state.get("rag_decision", {}),
        "policy_context_count": len(state.get("policy_results", {}).get("results", [])),
        "similar_exception_count": len(
            state.get("exception_results", {}).get("results", [])
        ),
    }
    return {
        **state,
        "denial_risk": risk,
        "route": route,
        "llm_decision": llm_decision,
        "final_summary": summary,
        "messages": _merge_messages(state, "Exception decision complete"),
    }


# ---------------- Scheduling graph nodes ----------------


def scheduling_parse_node(state: SchedulingState) -> SchedulingState:
    logger.info("Scheduling parse node")
    parsed = scheduling_request_parser.invoke({"request_text": state["request_text"]})
    extracted = parsed["extracted_request"]
    return {
        **state,
        "member_id": extracted.get("member_id") or state.get("member_id", ""),
        "extracted_request": extracted,
        "messages": _merge_messages(state, "Scheduling request parsed"),
    }


def scheduling_readiness_node(state: SchedulingState) -> SchedulingState:
    logger.info("Scheduling readiness node")
    member_id = state.get("member_id", "")
    specialty = state.get("extracted_request", {}).get("specialty")
    benefits = member_benefit_lookup.invoke(
        {"member_id": member_id, "specialty": specialty}
    )
    referrals = referral_lookup.invoke({"member_id": member_id, "specialty": specialty})

    auth_decision = _llm_json_decision(
        prompt=f"""
Decide whether this scheduling request needs an authorization lookup. Return JSON only.

Return this shape:
{{
  "check_authorization": true,
  "reason": "short explanation"
}}

Request:
{state.get('request_text')}

Extracted request:
{json.dumps(state.get('extracted_request'), indent=2)}

Benefit results:
{json.dumps(benefits, indent=2)[:3000]}

Referral results:
{json.dumps(referrals, indent=2)[:3000]}
""",
        fallback={
            "check_authorization": True,
            "reason": "Fallback rule: check authorization for payer scheduling readiness unless explicitly skipped by LLM.",
        },
    )

    if auth_decision.get("check_authorization"):
        auths = authorization_lookup.invoke(
            {"member_id": member_id, "specialty": specialty}
        )
    else:
        auths = {
            "ok": True,
            "skipped": True,
            "matches": [],
            "reason": auth_decision.get("reason"),
        }

    readiness = schedule_readiness_checker.invoke(
        {
            "benefit_results_json": json.dumps(benefits),
            "referral_results_json": json.dumps(referrals),
            "authorization_results_json": json.dumps(auths),
        }
    )

    return {
        **state,
        "benefit_results": benefits,
        "referral_results": referrals,
        "authorization_results": auths,
        "auth_decision": auth_decision,
        "schedule_readiness": readiness,
        "messages": _merge_messages(state, "Readiness checks complete"),
    }


def scheduling_readiness_router(state: SchedulingState) -> str:
    readiness = state.get("schedule_readiness", {})
    if readiness.get("ready_to_schedule"):
        logger.info("Scheduling router: ready, continue to provider matching")
        return "provider_match"
    logger.info("Scheduling router: not ready, finalize with human scheduler review")
    return "final"


def scheduling_provider_match_node(state: SchedulingState) -> SchedulingState:
    logger.info("Scheduling provider match node")
    extracted = state.get("extracted_request", {})
    specialty = extracted.get("specialty") or ""
    member = state.get("benefit_results", {}).get("member") or {}
    benefit_rows = state.get("benefit_results", {}).get("matches") or []
    plan_id = member.get("plan_id") or (
        benefit_rows[0].get("plan_id") if benefit_rows else ""
    )
    city = extracted.get("city") or ""

    query = json.dumps(
        {
            "request": extracted,
            "benefits": state.get("benefit_results"),
        }
    )
    matches = provider_specialty_rag_retriever.invoke({"query": query, "k": 5})

    locations = specialist_location_lookup.invoke(
        {"specialty": specialty, "city": city, "plan_id": plan_id}
    )

    provider_ranking = _llm_json_decision(
        prompt=f"""
Rank provider matches for this scheduling request. Return JSON only.

Return this shape:
{{
  "top_provider_ids_or_names": [],
  "ranking_reason": "short explanation",
  "escalate_to_human": false
}}

Request:
{state.get('request_text')}

Provider RAG matches:
{json.dumps(matches, indent=2)[:6000]}

Location matches:
{json.dumps(locations, indent=2)[:4000]}
""",
        fallback={
            "top_provider_ids_or_names": [],
            "ranking_reason": "LLM disabled or unavailable; use retrieved provider/location matches as-is.",
            "escalate_to_human": False,
        },
    )

    return {
        **state,
        "provider_matches": matches,
        "specialist_locations": locations,
        "provider_ranking": provider_ranking,
        "messages": _merge_messages(state, "Provider matching complete"),
    }


def scheduling_availability_node(state: SchedulingState) -> SchedulingState:
    logger.info("Scheduling availability node")
    extracted = state.get("extracted_request", {})
    specialty = extracted.get("specialty") or ""
    member = state.get("benefit_results", {}).get("member") or {}
    benefit_rows = state.get("benefit_results", {}).get("matches") or []
    plan_id = member.get("plan_id") or (
        benefit_rows[0].get("plan_id") if benefit_rows else ""
    )
    city = extracted.get("city") or ""

    availability = provider_availability_lookup.invoke(
        {"specialty": specialty, "plan_id": plan_id, "city": city}
    )
    if not availability.get("matches") and city:
        logger.info(
            f"No slots in preferred city={city}; expanding search to all matching locations"
        )
        availability = provider_availability_lookup.invoke(
            {"specialty": specialty, "plan_id": plan_id}
        )
    ranked = appointment_option_ranker.invoke(
        {
            "provider_matches_json": json.dumps(state.get("provider_matches", {})),
            "availability_json": json.dumps(availability),
        }
    )
    return {
        **state,
        "availability_results": availability,
        "appointment_options": ranked.get("appointment_options", []),
        "messages": _merge_messages(state, "Availability lookup complete"),
    }


def scheduling_final_node(state: SchedulingState) -> SchedulingState:
    logger.info("Scheduling final node")
    summary = scheduling_summary_writer.invoke(
        {
            "options_json": json.dumps(
                {"appointment_options": state.get("appointment_options", [])}
            )
        }
    )

    appointment_options = state.get("appointment_options", [])
    llm_summary = _llm_json_decision(
        prompt=f"""
Create the final scheduling decision summary. Return JSON only.

Return this shape:
{{
  "recommended_action": "offer_appointment_options|human_scheduler_review",
  "member_facing_summary": "brief explanation",
  "admin_notes": [],
  "appointment_options": []
}}

Request:
{state.get('request_text')}

Readiness:
{json.dumps(state.get('schedule_readiness', {}), indent=2)}

Provider ranking:
{json.dumps(state.get('provider_ranking', {}), indent=2)}

Appointment options:
{json.dumps(state.get('appointment_options', []), indent=2)}
""",
        fallback={
            **summary.get("final_summary", {}),
            "appointment_options": appointment_options,
            "member_facing_summary": (
                f"Found {len(appointment_options)} in-network appointment option(s) matching the request."
                if appointment_options
                else "Readiness checks completed; route to human scheduler review."
            ),
            "admin_notes": state.get("schedule_readiness", {}).get("issues", []),
        },
    )

    if appointment_options:
        llm_summary["appointment_options"] = appointment_options
        if not llm_summary.get("recommended_action"):
            llm_summary["recommended_action"] = "offer_appointment_options"

    return {
        **state,
        "final_summary": llm_summary,
        "messages": _merge_messages(state, "Scheduling summary complete"),
    }