"""Optional durable correction dispatches for early semantic model phases."""
import copy
import contextvars
import hashlib
import re
import tomllib
from dataclasses import dataclass, replace
from pathlib import Path
from importlib.resources import files

from viano_bug_analysis.case.contracts import read_json, validate, require_file
from viano_bug_analysis.lifecycle import State
from viano_bug_analysis.lifecycle.checkpoints import checkpoint_chain
from viano_bug_analysis.lifecycle.storage import json_bytes, key
from viano_bug_analysis.package import PackageReader
from .errors import InterpretationError

PHASES = {
    "interpretation": (State.INTERPRETING_CASE, "analysis/interpretations", "case-interpretation-rejected", "case-interpretation-error", "interpretation_context", "case-interpretation-context"),
    "hypotheses": (State.GENERATING_HYPOTHESES, "analysis/hypotheses", "hypotheses-rejected", "hypothesis-error", "hypothesis_context", "hypothesis-context"),
}
MODEL_CODES = frozenset({"LLM_OUTPUT_SCHEMA", "LLM_RESPONSE", "LLM_INCOMPLETE_OUTPUT",
    "INTERPRETATION_ENTITY_NAME_SOURCE", "INTERPRETATION_UNGROUNDED_ENTITY_LABEL", "INTERPRETATION_SCHEMA", "INTERPRETATION_SOURCE_UNKNOWN", "INTERPRETATION_SOURCE_KIND",
    "INTERPRETATION_EXPECTATION_UNKNOWN", "INTERPRETATION_CONTRADICTION_SOURCES",
    "INTERPRETATION_GAP_SOURCES", "INTERPRETATION_ENTITY", "INTERPRETATION_RESOLVED_ENTITY_OMITTED",
    "INTERPRETATION_UNRESOLVED_ENTITY_OMITTED", "INTERPRETATION_ENTITY_PROVENANCE", "INTERPRETATION_DESCRIPTION_SOURCES",
    "INTERPRETATION_PROVISIONAL_GAP_SOURCES", "HYPOTHESIS_SCHEMA", "HYPOTHESIS_COUNT",
    "HYPOTHESIS_DUPLICATE", "HYPOTHESIS_SOURCE_UNKNOWN", "HYPOTHESIS_GROUNDING", "HYPOTHESIS_PRIMARY_FOCUS"})
_attempt_context = contextvars.ContextVar("vba_early_correction", default=None)


def current_attempt(owner):
    value = _attempt_context.get()
    return value[1] if value and value[0] is owner else None


@dataclass(frozen=True)
class CorrectionConfig:
    max_invalid_responses: int = 3

    def __post_init__(self):
        if type(self.max_invalid_responses) is not int or self.max_invalid_responses <= 0:
            raise InterpretationError("CORRECTION_CONFIG")

    @classmethod
    def load(cls, path):
        try:
            with Path(path).open("rb") as stream:
                d = tomllib.load(stream)
            if set(d) != {"correction"}:
                raise ValueError()
            return cls(**d["correction"])
        except (OSError, ValueError, TypeError):
            raise InterpretationError("CORRECTION_CONFIG") from None


def digest(data):
    return hashlib.sha256(json_bytes(data)).hexdigest()


def base(phase, revision):
    return f"{PHASES[phase][1]}/{revision:04d}/corrections"


def identity(m, phase, kind):
    return {"format": "viano-early-correction", "format_version": 1, "kind": kind,
            "phase": phase, **{f: m[f] for f in ("case_id", "package_id", "analysis_revision")}}


def artifacts(reader, phase):
    result = []
    for _, cp in reversed(checkpoint_chain(reader)):
        d = cp["details"]
        if d.get("operation") == "early-correction" and d.get("phase") == phase and cp["analysis_revision"] == reader.manifest["analysis_revision"]:
            doc = read_json(reader, d["artifact_ref"], "early-correction", 4 * 1024**2)
            expected = base(phase, cp["analysis_revision"]) + "/" + (
                f"attempt-{doc.get('attempt', 0):04d}.json" if doc["kind"] == "attempt" else doc["kind"] + ".json")
            if (d["artifact_ref"] != expected or digest(doc) != d.get("artifact_sha256")
                    or any(doc[f] != reader.manifest[f] for f in ("case_id", "package_id", "analysis_revision"))
                    or doc["phase"] != phase):
                raise InterpretationError("CORRECTION_HISTORY")
            if doc["kind"] == "policy" and doc["feedback_prompt_version"] != "correction_v1":
                raise InterpretationError("CORRECTION_HISTORY")
            expected_fields = set(identity(reader.manifest, phase, doc["kind"])) | {
                "policy": {"max_invalid_responses", "feedback_prompt_version"},
                "attempt": {"attempt", "attempt_id", "previous_error_ref", "feedback"},
                "failure": {"policy", "error_refs", "reason", "context_ref", "context_sha256", "remaining_questions", "recommended_human_actions"},
            }[doc["kind"]]
            if set(doc) != expected_fields:
                raise InterpretationError("CORRECTION_HISTORY")
            result.append((d["artifact_ref"], doc))
    return result


def failures(reader, phase):
    result, calls = [], set()
    for _, cp in reversed(checkpoint_chain(reader)):
        d = cp["details"]
        if cp["analysis_revision"] != reader.manifest["analysis_revision"] or d.get("operation") != PHASES[phase][2] or d.get("code") not in MODEL_CODES:
            continue
        error = read_json(reader, d["error_ref"], PHASES[phase][3], 4 * 1024**2)
        call = read_json(reader, d["llm_result_ref"], "llm-call", 4 * 1024**2)
        if (error["code"] != d["code"] or error["llm_result_ref"] != d["llm_result_ref"]
                or call["phase"] != PHASES[phase][0] or call["status"] not in {"SUCCESS", "ERROR"}
                or any(call[f] != reader.manifest[f] or error[f] != reader.manifest[f] for f in ("case_id", "package_id", "analysis_revision"))):
            raise InterpretationError("CORRECTION_HISTORY")
        context = read_json(reader, error["context_ref"], PHASES[phase][5], 4 * 1024**2)
        require_file(reader, call["request_ref"])
        with reader._archive.open(call["request_ref"]) as stream:
            from viano_bug_analysis.package.manifest import parse_manifest
            request = parse_manifest(reader.validator._read_bounded(stream, min(4 * 1024**2, reader.validator.limits.max_file_bytes)))
        from .prompt import InterpretationPromptBuilder
        from .hypothesis_prompt import HypothesisPromptBuilder
        expected = (InterpretationPromptBuilder() if phase == "interpretation" else HypothesisPromptBuilder()).build(context)
        if (error["context_sha256"] != digest(context) or digest(request) != call["input_request_hash"]
                or call["result_ref"] != d["llm_result_ref"] or request.get("messages", [])[:2] != expected["messages"]
                or request.get("output_schema") != expected["output_schema"]):
            raise InterpretationError("CORRECTION_HISTORY")
        if call["status"] == "SUCCESS":
            from .contracts import validate_interpretation
            from .hypothesis_contracts import validate_proposals
            from .hypothesis_config import HypothesisError
            try:
                (validate_interpretation if phase == "interpretation" else validate_proposals)(call["outcome"]["document"], context)
            except (InterpretationError, HypothesisError) as invalid:
                if invalid.code != error["code"]:
                    raise InterpretationError("CORRECTION_HISTORY")
            else:
                raise InterpretationError("CORRECTION_HISTORY")
        elif call["code"] != error["code"] or call["code"] not in MODEL_CODES:
            raise InterpretationError("CORRECTION_HISTORY")
        if d["llm_result_ref"] not in calls:
            result.append((d["error_ref"], error, call))
            calls.add(d["llm_result_ref"])
    return result


def feedback(reader, phase, failure):
    ref, error, call = failure
    context = read_json(reader, error["context_ref"], PHASES[phase][5], 4 * 1024**2)
    payload = call["outcome"].get("document")
    issues = []
    if phase == "interpretation" and isinstance(payload, dict):
        allowed_fields = {"observed_behavior": {"observation", "report", "tester_comment", "reproduction", "step"},
                          "expected_behavior": {"expectation"}, "tester_assumptions": {"tester_assumption", "tester_comment"},
                          "confirmed_technical_facts": {"fact"}}
        for field, allowed in allowed_fields.items():
            for index, statement in enumerate(payload.get(field, [])):
                if not isinstance(statement, dict):
                    continue
                for sid in statement.get("source_ids", []):
                    if not isinstance(sid, str) or not re.fullmatch(r"[A-Z][0-9]{4,}", sid):
                        continue
                    actual = context["sources"].get(sid, {}).get("kind")
                    if actual not in allowed:
                        issue = {"field": f"{field}/{index}", "source_id": sid,
                                 "actual_kind": actual, "allowed_kinds": sorted(allowed)}
                        available = sorted(s for s, entry in context["sources"].items() if entry["kind"] in allowed)
                        if len(available) <= 32:
                            issue["allowed_source_ids"] = available
                        issues.append(issue)
        if error.get('prompt_version') in {'case_interpretation_resolved_v2','case_interpretation_resolved_v3','case_interpretation_resolved_v4','case_interpretation_resolved_v5','case_interpretation_resolved_v6'}:
            from .contracts import entity_citation_issues
            issues.extend(entity_citation_issues(payload,context))
            if error.get('prompt_version') in {'case_interpretation_resolved_v5','case_interpretation_resolved_v6'}:
                for index,entity in enumerate(payload.get('relevant_entities',[])):
                    source = context['sources'].get(entity.get('entity_id'),{})
                    bindings = source.get('data',{}).get('name_bindings',[])
                    allowed = sorted({b['source_id'] for b in bindings})
                    if allowed and not set(entity.get('source_ids',[])) & set(allowed):
                        issues.append({'field':f'relevant_entities/{index}/source_ids',
                            'constraint':'cite_typed_name_binding','allowed_source_ids':allowed})
            if error.get('prompt_version') in {'case_interpretation_resolved_v4','case_interpretation_resolved_v5','case_interpretation_resolved_v6'}:
                from .quality import unnamed_entity_reasons
                reasons = unnamed_entity_reasons(context)
                for index, entity in enumerate(payload.get('relevant_entities',[])):
                    sid = entity.get('entity_id')
                    if sid in reasons and entity.get('reason') != reasons[sid]:
                        issues.append({'field':f'relevant_entities/{index}/reason',
                            'constraint':'use_neutral_identity_without_name_binding', 'required_reason':reasons[sid]})
            for sid, source in context['sources'].items():
                if (source['kind']=='data_quality' and source['case_pointer'].startswith('/entity_resolution/')
                        and not any(sid in s.get('source_ids',[]) for s in payload.get('unknowns',[]) if isinstance(s,dict))):
                    issues.append({'field':'unknowns','source_id':sid,'constraint':'retain_scoped_lookup_uncertainty'})
            if error.get('prompt_version') in {'case_interpretation_resolved_v3','case_interpretation_resolved_v4','case_interpretation_resolved_v5','case_interpretation_resolved_v6'}:
                from .quality import description_sources
                allowed = description_sources(context)
                for index, statement in enumerate(payload.get('described_behavior',[])):
                    if not set(statement.get('source_ids',[])) <= allowed:
                        issues.append({'field':f'described_behavior/{index}/source_ids',
                                       'constraint':'cite_local_description','allowed_source_ids':sorted(allowed)[:32]})
                gap = payload.get('provisional_behavior_gap')
                described = {sid for s in payload.get('described_behavior',[]) for sid in s.get('source_ids',[])}
                gap_refs = set(gap.get('source_ids',[])) if isinstance(gap,dict) else set()
                observed = {sid for sid,s in context['sources'].items() if s['kind'] in {'observation','report','tester_comment'}}
                if gap is not None and (not gap_refs & described or not gap_refs & observed
                    or any(context['sources'].get(sid,{}).get('kind') not in {'observation','report','tester_comment','reproduction','step','fact'} for sid in gap_refs)):
                    issues.append({'field':'provisional_behavior_gap/source_ids',
                        'constraint':'cite_described_behavior_and_observation',
                        'description_source_ids':sorted(allowed)[:32],
                        'observation_source_ids':[sid for sid,s in context['sources'].items()
                                                  if s['kind'] in {'observation','report','tester_comment'}][:32]})
        issues = issues[:32]
    if phase == "hypotheses" and isinstance(payload, dict):
        from .hypothesis_contracts import _normalized
        titles, mechanisms = {}, {}
        proposals = payload.get("hypotheses", [])
        if not context["policy"]["min_hypotheses"] <= len(proposals) <= context["policy"]["max_hypotheses"]:
            issues.append({"field": "hypotheses", "constraint": "count", "minimum": context["policy"]["min_hypotheses"], "maximum": context["policy"]["max_hypotheses"]})
        for index, proposal in enumerate(proposals):
            if not isinstance(proposal, dict): continue
            for field, seen in (("title", titles), ("mechanism", mechanisms)):
                value = _normalized(proposal.get(field, ""))
                if value in seen:
                    issues.append({"field": f"hypotheses/{index}/{field}", "constraint": "distinct", "conflicts_with": f"hypotheses/{seen[value]}/{field}"})
                seen[value] = index
            for field, sources in (("source_ids", context["case_context"]["sources"]), ("interpretation_ids", context["interpretation_sources"])):
                for sid in proposal.get(field, []):
                    if isinstance(sid, str) and re.fullmatch(r"[A-Z][0-9]{4,}", sid) and sid not in sources:
                        issues.append({"field": f"hypotheses/{index}/{field}", "source_id": sid, "constraint": "known_source"})
            kinds = {context["case_context"]["sources"].get(sid, {}).get("kind") for sid in proposal.get("source_ids", [])}
            if not kinds & {"observation", "report", "tester_comment", "fact"}:
                issues.append({"field": f"hypotheses/{index}/source_ids", "constraint": "cite_observation_or_inline_fact"})
            if index == 0:
                focus = set()
                for kind in ("first_known_divergence", "failure_stage", "observed_behavior"):
                    focus = {sid for sid, source in context["interpretation_sources"].items() if source["kind"] == kind}
                    if focus: break
                if not set(proposal.get("interpretation_ids", [])) & focus:
                    issues.append({"field": "hypotheses/0/interpretation_ids", "constraint": "primary_focus", "allowed_source_ids": sorted(focus)[:32]})
    instruction = "Return a complete replacement JSON object conforming to the supplied schema and source rules. Do not invent facts or sources. Use empty arrays/null for unsupported claims. Distinct hypotheses must retain valid grounding and count constraints."
    if phase == 'interpretation' and error.get('prompt_version') in {'case_interpretation_resolved_v2','case_interpretation_resolved_v3','case_interpretation_resolved_v4','case_interpretation_resolved_v5','case_interpretation_resolved_v6'}:
        instruction += " Keep tester observations in observed_behavior with observation sources; move technical assertions to confirmed_technical_facts with fact sources. Retain selected entity_id values and cite their original mentions plus entity or matching capture evidence. Keep lookup uncertainties scoped to their queries."
    return {"code":error["code"],"issues":issues[:32],"instruction":instruction}


def augment(request, reader, phase, attempt):
    result = copy.deepcopy(request)
    if attempt.get("previous_error_ref"):
        result["messages"].append({"role": "system", "content": files("viano_bug_analysis.llm.prompts").joinpath("correction_v1.txt").read_text(encoding="utf-8")})
        history = failures(reader, phase)
        failure = next((f for f in history if f[0] == attempt["previous_error_ref"]), None)
        if failure is None or feedback(reader, phase, failure) != attempt["feedback"]:
            raise InterpretationError("CORRECTION_FEEDBACK_CHANGED")
        payload = failure[2]["outcome"].get("document")
        if payload is not None:
            result["messages"].append({"role": "assistant", "content": json_bytes(payload).decode()})
        result["messages"].append({"role": "user", "content": json_bytes({"task": "correct_rejected_output", "validation_feedback": attempt["feedback"]}).decode()})
    return result


def replay_request(reader, request, call, phase):
    for _, doc in artifacts(reader, phase):
        if doc["kind"] == "attempt":
            operation = ("case-interpretation:v1:" if phase == "interpretation" else "hypotheses:v1:") + doc["attempt_id"]
            if call["operation_key"] == key(f"{reader.manifest['analysis_revision']}:{operation}"):
                return augment(request, reader, phase, doc)
    return request


def save(owner, m, phase, doc):
    validate(doc, "early-correction")
    ref = base(phase, m["analysis_revision"]) + "/" + (
        f"attempt-{doc['attempt']:04d}.json" if doc["kind"] == "attempt" else doc["kind"] + ".json")
    owner.store.advance(m["case_id"], m["package_id"], PHASES[phase][0],
        step_id=f"correction:{phase}:{m['analysis_revision']}:{doc['kind']}:{doc.get('attempt', 0)}",
        details={"operation": "early-correction", "phase": phase, "artifact_ref": ref, "artifact_sha256": digest(doc)},
        additions={ref: json_bytes(doc)}, expected_generation=m["package_generation"])
    return ref


def review(owner, reader, phase, policy, rejected):
    m = reader.manifest
    doc = {**identity(m, phase, "failure"), "policy": policy,
           "error_refs": [f[0] for f in rejected], "reason": "MODEL_VALIDATION_EXHAUSTED",
           "context_ref": m["artifacts"].get(PHASES[phase][4]),
           "remaining_questions": ["No fully validated model output was accepted for this phase."],
           "recommended_human_actions": ["Inspect the rejected outputs and validation feedback; supply corrected observations or expected behavior before reopening the case."]}
    doc["context_sha256"] = digest(read_json(reader, doc["context_ref"], PHASES[phase][5], 4 * 1024**2))
    ref = base(phase, m["analysis_revision"]) + "/failure.json"
    validate(doc, "early-correction")
    path = owner.store.advance(m["case_id"], m["package_id"], State.FAILED_TO_ANALYZE,
        step_id=f"correction:{phase}:{m['analysis_revision']}:review",
        reason="MODEL_VALIDATION_EXHAUSTED", details={"operation": "early-correction", "phase": phase,
            "artifact_ref": ref, "artifact_sha256": digest(doc)}, additions={ref: json_bytes(doc)},
        artifacts={"failure_report": ref, "change_request": None, "change_request_summary": None},
        expected_generation=m["package_generation"])
    return path


async def dispatch(owner, phase, case_id, package_id, attempt_id, callback):
    latest = owner.store.latest(case_id, package_id)
    if latest is None:
        return await callback(case_id, package_id, attempt_id)
    with PackageReader(latest, owner.store.limits) as reader:
        m = reader.manifest
        existing = artifacts(reader, phase)
        policy = next((d for _, d in existing if d["kind"] == "policy"), None)
        complete = "case-interpreter" if phase == "interpretation" else "hypothesis-builder"
        completed = any(cp["step_id"] == f"{complete}:v1:revision:{m['analysis_revision']}:complete" for _, cp in checkpoint_chain(reader))
        if completed or m["state"] == State.BLOCKED or (policy is None and owner.correction_config is None):
            return await callback(case_id, package_id, attempt_id)
        if m["state"] == State.FAILED_TO_ANALYZE:
            validate_review(reader)
            return owner._correction_result("INCONCLUSIVE", latest, m["analysis_revision"])
        if owner.correction_config is not None and policy and policy["max_invalid_responses"] != owner.correction_config.max_invalid_responses:
            raise InterpretationError("CORRECTION_POLICY_CHANGED")
        rejected = failures(reader, phase)
        if policy is None:
            policy = {**identity(m, phase, "policy"), "max_invalid_responses": owner.correction_config.max_invalid_responses,
                      "feedback_prompt_version": "correction_v1"}
            save(owner, m, phase, policy)
    with PackageReader(owner.store.latest(case_id, package_id), owner.store.limits) as reader:
        m = reader.manifest
        if len(rejected) >= policy["max_invalid_responses"]:
            path = review(owner, reader, phase, policy, rejected)
            return owner._correction_result("INCONCLUSIVE", path, m["analysis_revision"])
        number = len(rejected) + 1
        attempt = {**identity(m, phase, "attempt"), "attempt": number,
                   "attempt_id": f"correction-{phase}-{number}", "previous_error_ref": rejected[-1][0] if rejected else None,
                   "feedback": feedback(reader, phase, rejected[-1]) if rejected else None}
        existing_attempt = next((d for _, d in artifacts(reader, phase) if d["kind"] == "attempt" and d["attempt"] == number), None)
        if existing_attempt and existing_attempt != attempt:
            raise InterpretationError("CORRECTION_HISTORY")
        if not existing_attempt:
            save(owner, m, phase, attempt)
    token = _attempt_context.set((owner, attempt))
    try:
        result = await callback(case_id, package_id, attempt["attempt_id"])
    finally:
        _attempt_context.reset(token)
    if result.status == "ERROR" and result.code in MODEL_CODES:
        with PackageReader(result.package, owner.store.limits) as reader:
            rejected = failures(reader, phase)
            if len(rejected) >= policy["max_invalid_responses"]:
                path = review(owner, reader, phase, policy, rejected)
                return owner._correction_result("INCONCLUSIVE", path, reader.manifest["analysis_revision"])
        return replace(result, status="REJECTED")
    return result


def validate_review(reader):
    m = reader.manifest
    documents = [(p, d) for phase in PHASES for p, d in artifacts(reader, phase) if d["kind"] == "failure"]
    if len(documents) != 1:
        raise InterpretationError("CORRECTION_REVIEW")
    ref, report = documents[0]
    phase = report["phase"]
    policy = next((d for _, d in artifacts(reader, phase) if d["kind"] == "policy"), None)
    rejected = failures(reader, phase)
    if (policy != report["policy"] or len(rejected) < policy["max_invalid_responses"]
            or report["error_refs"] != [f[0] for f in rejected]
            or report["context_ref"] != m["artifacts"].get(PHASES[phase][4])
            or report["context_sha256"] != digest(read_json(reader, report["context_ref"], PHASES[phase][5], 4 * 1024**2))
            or m["artifacts"].get("failure_report") != ref or m["artifacts"].get("change_request")
            or m["artifacts"].get("diagnosis") or m["artifacts"].get("hypothesis_registry")
            or (phase == "interpretation" and m["artifacts"].get("case_interpretation"))
            or m["state"] != State.FAILED_TO_ANALYZE or m["routing"]["destination"] != "review"
            or report["reason"] != "MODEL_VALIDATION_EXHAUSTED"
            or report["remaining_questions"] != ["No fully validated model output was accepted for this phase."]
            or report["recommended_human_actions"] != ["Inspect the rejected outputs and validation feedback; supply corrected observations or expected behavior before reopening the case."]):
        raise InterpretationError("CORRECTION_REVIEW")
    return report