"""
Task 2: Vulnerability Remediation (Medium)

Agent must propose correct fixes for identified vulnerabilities.
Scored on: correct fix version, valid replacement code, budget management.
"""
from __future__ import annotations

from typing import Any, Dict, Tuple

from .base import TaskHandler, apply_identification, apply_remediation
from env.models import Action, EngineState
from env.verification import task_completion_score


class Task2Handler(TaskHandler):
    """
    Remediation task: agent proposes fixes for code vulnerabilities.

    Actions:
    - identify: find vulns first (same as task 1)
    - remediate: propose a fix for a specific CVE
    - done: end episode

    Scoring: fraction of vulns correctly remediated, minus budget overruns.
    """

    def execute(self, state: EngineState, action: Action) -> Tuple[float, Dict[str, Any]]:
        if action.action_type == "done":
            state.done = True
            sc = task_completion_score(state, state.task_id)
            return sc, {"sc": round(sc, 4)}

        if action.action_type == "identify":
            return self._handle_identify(state, action)

        if action.action_type == "remediate":
            return self._handle_remediate(state, action)

        state.last_action_error = f"Unsupported: {action.action_type}"
        return 0.0, {}

    def _handle_identify(self, state: EngineState, action: Action) -> Tuple[float, Dict[str, Any]]:
        if not action.findings:
            state.last_action_error = "No findings"
            return 0.0, {}

        info = apply_identification(state, action.findings)
        return task_completion_score(state, state.task_id), {
            "identified": info["accepted"],
            "partial": info["partial"],
            "fp": info["fp"],
            "evidence_gain": info["evidence_gain"],
        }

    def _handle_remediate(self, state: EngineState, action: Action) -> Tuple[float, Dict[str, Any]]:
        if not action.remediation:
            state.last_action_error = "No remediation provided"
            return 0.0, {}

        rem = action.remediation
        cve_id = rem.cve_id

        if cve_id not in state.ground_truth_vulns:
            state.invalid_remediations += 1
            state.last_action_error = f"CVE {cve_id} not in ground truth"
            return -0.15, {"error": "not_present"}

        if cve_id in state.remediated_vulns:
            state.last_action_error = "Already remediated"
            return -0.05, {}

        if state.budget_points < 2:
            state.invalid_remediations += 1
            state.last_action_error = "Insufficient budget"
            return -0.1, {}
        state.budget_points -= 2

        info = {"cve": cve_id, **apply_remediation(state, rem)}

        all_fixed = len(state.remediated_vulns) >= state.initial_vuln_count
        if all_fixed:
            state.done = True
            sc = task_completion_score(state, state.task_id)
            return sc, {**info, "sc": round(sc, 4)}

        return task_completion_score(state, state.task_id), info