from __future__ import annotations
from typing import Optional
from codestar.core.schemas import ProblemMeta, SubmissionRequest, JudgeResult, JudgeStatus
from codestar.core.ast_checker import check_code_ast
from codestar.core.runner import execute_submission


class JudgeEngine:
    def __init__(self):
        pass

    def judge(
        self,
        problem: ProblemMeta,
        user_code: str,
        custom_testcase: Optional[str] = None,
    ) -> JudgeResult:
        # Step 1: AST Check (Syntax & Security)
        is_safe, error_msg, ast_status = check_code_ast(
            code=user_code,
            extra_forbidden_modules=problem.forbidden_modules,
            extra_forbidden_calls=problem.forbidden_calls,
        )
        if not is_safe:
            return JudgeResult(
                status=ast_status or JudgeStatus.CE,
                message=error_msg,
                passed_count=0,
                total_count=0,
                execution_time_ms=0.0,
            )

        # Step 1b: AST Check on custom_testcase if provided
        if custom_testcase and custom_testcase.strip():
            c_safe, c_msg, c_status = check_code_ast(
                code=custom_testcase,
                extra_forbidden_modules=problem.forbidden_modules,
                extra_forbidden_calls=problem.forbidden_calls,
            )
            if not c_safe:
                return JudgeResult(
                    status=c_status or JudgeStatus.CE,
                    message=f"Custom testcase error: {c_msg}",
                    passed_count=0,
                    total_count=0,
                    execution_time_ms=0.0,
                )

        # Step 2: Execute in Runner Harness
        return execute_submission(problem, user_code, custom_testcase=custom_testcase)