"""Result formatter for Tru-Guard - formats results for TRUSYS backend."""

import json
import time
from typing import Any

from trusys.types import AggregatedResult, ScanContext, ScanType


class ResultFormatter:
    """Formats guardrail results for TRUSYS backend storage."""

    @staticmethod
    def format_for_backend(
        result: AggregatedResult,
        context: ScanContext,
        application_id: str | None = None,
        application_version: str | None = None,
        environment: str | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        """Format aggregated result for TRUSYS backend.

        When a scanner fails due to an error (e.g. HTTP error), component_results,
        pass, and score are set to null and the error message is set in the error field.
        When scanners complete normally (with or without violations), component_results,
        pass, and score are populated and error is null.

        Args:
            result: Aggregated result from guardrail execution
            context: Scan context containing prompt and response
            application_id: Application ID (optional, can be from config)

        Returns:
            Dictionary in the format expected by TRUSYS backend
        """
        # Check if any scanner failed due to an error (e.g. HTTP, timeout) - metadata["error"] is set
        scanner_errors = [
            scan_result.metadata.get("error")
            for scan_result in result.results
            if scan_result.metadata.get("error")
        ]

        if scanner_errors:
            # Scanner failure: component_results, pass, score are null; error holds the message(s)
            error_message = "; ".join(str(e) for e in scanner_errors)
            formatted_result = {
                "application_id": application_id or "",
                "app_version": application_version,
                "environment": environment,
                "guardrail_stage": context.scan_type.value,  # "input" or "output"
                "prompt": context.prompt or "",
                "response": context.response or "",
                "final_content": result.final_content,
                "latency": result.total_execution_time_ms,  # Keep in ms
                "pass": None,
                "score": None,
                "error": error_message,
                "execution_time": time.time(),  # Current timestamp
                "componentResults": None,
                "metadata": metadata or {},
            }
            return formatted_result

        # Normal result: build component_results, pass, score; error is null
        scores = []
        for scan_result in result.results:
            score = scan_result.metadata.get("score")
            if score is not None:
                scores.append(score)

        if scores:
            overall_score = sum(scores) / len(scores)
        else:
            overall_score = 1.0 if result.passed else 0.0

        # Map scanner_name -> ActionRecord for enriching action with original/modified content
        actions_by_scanner = {r.scanner_name: r for r in result.actions_taken}

        component_results = []
        for scan_result in result.results:
            action = dict(scan_result.metadata.get("action", {"type": "ignore"}))
            action_record = actions_by_scanner.get(scan_result.scanner_name)
            if action_record is not None:
                action["original_content"] = action_record.original_content
                action["modified_content"] = action_record.modified_content

            component_result = {
                "assertion": scan_result.metadata.get("assertion", {
                    "type": scan_result.scanner_name,
                    "output": []
                }),
                "pass": scan_result.passed,
                "reason": scan_result.metadata.get("reason", ""),
                "score": scan_result.metadata.get("score", 1.0 if scan_result.passed else 0.0),
                "latency": scan_result.metadata.get("latency", scan_result.execution_time_ms),
                "action": action,
            }
            if "tokensUsed" in scan_result.metadata:
                component_result["tokensUsed"] = scan_result.metadata["tokensUsed"]
            component_results.append(component_result)

        formatted_result = {
            "application_id": application_id or "",
            "app_version": application_version,
            "environment": environment,
            "guardrail_stage": context.scan_type.value,  # "input" or "output"
            "prompt": context.prompt or "",
            "response": context.response or "",
            "final_content": result.final_content,
            "latency": result.total_execution_time_ms,  # Keep in ms
            "pass": result.passed,
            "score": overall_score,
            "error": None,
            "execution_time": time.time(),  # Current timestamp
            "componentResults": component_results,
            "metadata": metadata or {},
        }

        return formatted_result

    @staticmethod
    def print_formatted_result(
        result: AggregatedResult,
        context: ScanContext,
        application_id: str | None = None,
        application_version: str | None = None,
        environment: str | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> None:
        """Format and print result to console (for debugging/development).

        Args:
            result: Aggregated result from guardrail execution
            context: Scan context containing prompt and response
            application_id: Application ID (optional)
            application_version: Application version (optional)
            environment: Deployment environment (optional)
            metadata: Additional metadata to attach to the result
        """
        formatted = ResultFormatter.format_for_backend(
            result,
            context,
            application_id=application_id,
            application_version=application_version,
            environment=environment,
            metadata=metadata,
        )
        print("\n" + "=" * 80)
        print("GUARDRAIL RESULT")
        print("=" * 80)
        print(json.dumps(formatted, indent=2, ensure_ascii=False))
        print("=" * 80 + "\n")
