# Copyright UL Research Institutes
# SPDX-License-Identifier: Apache-2.0

import ast
import re
from typing import Iterable, Optional

import javalang
import pyarrow
import pyarrow.dataset
import pydantic

from dyff.audit.scoring.base import Rubric
from dyff.schema.base import int32
from dyff.schema.dataset import ReplicatedItem
from dyff.schema.dataset.arrow import arrow_schema, schema_function


def is_empty_python_code(code):
    code = str(code).strip()  # whitespaces
    code = re.sub(r"#.*", "", code)  # comments
    code = re.sub(r"(\'[\s\S]*?\'|\"[\s\S]*?\")", "", code)  # docstrings
    code = str(code).strip()  # whitespaces
    return not code


def is_valid_python_syntax(code):
    if not code or is_empty_python_code(code):
        return False

    try:
        tree = ast.parse(code)
        return tree is not None
    except:
        return False


def is_empty_java_code(code):
    code = str(code).strip()  # whitespaces
    code = re.sub(r"//.*", "", code)  # single-line comment
    code = re.sub(
        r"/\*.*?\*/", "", code, flags=re.DOTALL | re.MULTILINE
    )  # multilines or javadoc
    code = str(code).strip()  # whitespaces
    return not code


def is_valid_java_syntax(code):
    if not code or is_empty_java_code(code):
        return False

    tree = None
    try:
        tree = javalang.parse.parse(code)
    except:
        code = "class T {\n" + code + "\n}"
        try:
            tree = javalang.parse.parse(code)
        except:
            pass
    return tree is not None


class GetValidSyntaxScoreItem(ReplicatedItem):
    language: str = pydantic.Field(description="The target programming language")
    valid: Optional[int32(ge=0, le=1)] = pydantic.Field(  # type: ignore
        default=None,
        description="Valid syntax score for a code generated by the LLM",
    )


class GetValidSyntaxScore(Rubric):
    """
    Get the syntax validation score for code generated by an LLM.
    Return 0 for invalid syntax and 1 for valid syntax.
    """

    @property
    def name(self) -> str:
        return "evaluator_code_rubric.GetValidSyntaxScore"

    @property
    @schema_function(arrow_schema(GetValidSyntaxScoreItem))
    def schema(self) -> pyarrow.Schema:
        pass

    def apply(
        self, input_data: pyarrow.dataset.Dataset, output_data: pyarrow.dataset.Dataset
    ) -> Iterable[pyarrow.RecordBatch]:
        input_items = {}
        for input_batch in input_data.to_batches(
            columns=["_index_", "language", "text"]
        ):
            for input_item in input_batch.to_pylist():
                input_items[input_item["_index_"]] = input_item

        for output_batch in output_data.to_batches(
            columns=["_index_", "_replication_", "responses"]
        ):
            result_batch = []
            for output_item in output_batch.to_pylist():
                index = output_item["_index_"]
                input_item = input_items[index]
                replication = output_item["_replication_"]
                language = input_item["language"]
                input_text = input_item["text"]
                output_text = output_item["responses"][0]["text"]
                try:
                    # get code from responses
                    output_code = output_text.replace(input_text, "")
                    output_code = str(output_code).split("</code>")[0]

                    is_valid = None
                    if str(language).lower() == "python":
                        is_valid = 1 if is_valid_python_syntax(output_code) else 0
                    elif str(language).lower() == "java":
                        is_valid = 1 if is_valid_java_syntax(output_code) else 0
                    result = GetValidSyntaxScoreItem(
                        _index_=index,
                        _replication_=replication,
                        language=language,
                        valid=is_valid,
                    ).dict()
                except:
                    result = GetValidSyntaxScoreItem(
                        _index_=index,
                        _replication_=replication,
                        language=language,
                        valid=None,
                    ).dict()
                result_batch.append(result)

            yield pyarrow.RecordBatch.from_pylist(result_batch, schema=self.schema)
