# Description: This file contains the reporting contract and utilities for
# rendering reports

import json
import typing
import proto.dtx.messages.common.threat_pb2 as threat_pb2

from abc import ABC, abstractmethod
from detoxio.models import LLMScanResult
from google.protobuf.json_format import MessageToDict

from detoxio.utils import get_threat_category_human_name

class LLMScanReport(ABC):
    """
    LLMScanReport is an abstract base class that defines the contract for
    rendering reports from LLM scans.
    """
    @abstractmethod
    def render(self, results: LLMScanResult):
        """
        Render the report from the results.
        """
        raise NotImplementedError("Subclasses must implement this method")

class JSONLinesLLMScanReport(LLMScanReport):
    """
    JSONLinesLLMScanReport renders each prompt evaluation result as JSON lines.
    This reporting format is suitable for tool integrations, where the tool expects
    raw results for its own rendering and visualization.
    """

    def __init__(self, file: typing.TextIO = None):
        if file is None:
            raise ValueError("file is required for JSONLinesLLMScanReport")

        self.file = file

    def render(self, results: LLMScanResult):
        for result in results:
            self.file.write(json.dumps(MessageToDict(result, preserving_proto_field_name=True)) + "\n")

class JupyterNotebookLLMScanReport(LLMScanReport):
    """
    [WORK IN PROGRESS]: Jupyter notebook reporter
    """
    def render(self, results: LLMScanResult):
        env = None
        try:
            env = self.jupyter_notebook_environment()
        except ImportError:
            raise NotImplementedError("This reporter is only supported in Jupyter Notebook environment")

        display = env['display'] # noqa: F841
        HTML = env['HTML'] # noqa: F841

    def jupyter_notebook_environment(self):
        from IPython.core.display import display, HTML # noqa: F401
        return locals()

class MarkdownLLMScanReport(LLMScanReport):
    def __init__(self, writer: typing.TextIO):
        self.writer = writer

    def render(self, results: LLMScanResult):
        self.writer.write("# LLM Scan Report\n\n")

        self.writer.write("## Summary\n")
        self.writer.write(f"Total Tests: **{results.total_count()}** Failed Tests: **{results.failed_count()}**\n")
        self.writer.write("\n")

        self.writer.write("### Threats Detected\n")
        threat_group = results.group_by_threats()
        for threat_class, threat_categories in threat_group.items():
            for threat_category, threats in threat_categories.items():
                self.writer.write(f"- {get_threat_category_human_name(threat_category)}: {len(threats)}\n")

        self.writer.write("\n## Proof of Concept\n")

        idx = 0
        for result in results.failed_results():
            self.writer.write(f"### POC/{idx}\n")

            self.writer.write("#### Prompt\n")
            self.writer.write(f"\n```\n{result.prompt.data.content}\n```\n")

            self.writer.write("#### Responses\n")
            for response in result.responses:
                self.writer.write(f"\n```\n{response.response.message.content}\n```\n")

                for vulnerability in response.results:
                    if vulnerability.status == threat_pb2.THREAT_EVALUATION_STATUS_SAFE:
                        continue
                    self.writer.write(f"- {get_threat_category_human_name(vulnerability.threat.threat_category)}\n")

            idx += 1

        self.writer.write("\n\n")
        self.writer.write("> Report generated by [detoxio.ai](https://docs.detoxio.ai)")
