from __future__ import annotations

import logging
from typing import Any, cast

from asgiref.sync import sync_to_async

from core.log_features.models import OptimizationFeatures
from core.shared.optimizer_schemas import OptimizeResponseItemSchema, OptimizeResponseSchema
from core.shared.testgen_models import TestGenResponseSchema

logger = logging.getLogger(__name__)


@sync_to_async
def get_rerun_record(rerun_trace_id: str) -> OptimizationFeatures | None:
    try:
        return cast(OptimizationFeatures, OptimizationFeatures.objects.get(trace_id=rerun_trace_id))
    except OptimizationFeatures.DoesNotExist:
        return None


def rerun_optimize(record: OptimizationFeatures, source_filter: str) -> OptimizeResponseSchema | None:
    origin: dict[str, Any] = cast("dict[str, Any]", record.optimizations_origin) or {}
    post: dict[str, Any] = cast("dict[str, Any]", record.optimizations_post) or {}
    explanations: dict[str, Any] = cast("dict[str, Any]", record.explanations_post) or {}

    items: list[OptimizeResponseItemSchema] = []
    for opt_id, meta in origin.items():
        if meta.get("source") != source_filter:
            continue
        code = post.get(opt_id)
        if code is None:
            continue
        items.append(
            OptimizeResponseItemSchema(
                source_code=code,
                explanation=explanations.get(opt_id, ""),
                optimization_id=opt_id,
                parent_id=meta.get("parent"),
            )
        )

    if not items:
        return None
    return OptimizeResponseSchema(optimizations=items)


def rerun_testgen(record: OptimizationFeatures, test_index: int) -> TestGenResponseSchema | None:
    generated: list[str | None] = cast("list[str | None]", record.generated_test) or []
    instrumented: list[str | None] = cast("list[str | None]", record.instrumented_generated_test) or []
    perf: list[str | None] = cast("list[str | None]", record.instrumented_perf_test) or []

    if test_index >= len(generated) or test_index >= len(instrumented):
        return None

    # Check if values at the index are None (can happen with NULL in database arrays)
    generated_val = generated[test_index]
    instrumented_val = instrumented[test_index]

    if generated_val is None or instrumented_val is None:
        return None

    perf_val = perf[test_index] if test_index < len(perf) else None
    # Default to empty string if perf value is None
    perf_val = perf_val if perf_val is not None else ""

    return TestGenResponseSchema(
        generated_tests=generated_val,
        instrumented_behavior_tests=instrumented_val,
        instrumented_perf_tests=perf_val,
        raw_generated_tests=None,
    )


def rerun_refinement(
    record: OptimizationFeatures, parent_optimization_ids: list[str]
) -> list[OptimizeResponseItemSchema]:
    origin: dict[str, Any] = cast("dict[str, Any]", record.optimizations_origin) or {}
    post: dict[str, Any] = cast("dict[str, Any]", record.optimizations_post) or {}
    explanations: dict[str, Any] = cast("dict[str, Any]", record.explanations_post) or {}
    parent_set = set(parent_optimization_ids)

    items: list[OptimizeResponseItemSchema] = []
    for opt_id, meta in origin.items():
        if meta.get("source") != "REFINE":
            continue
        if meta.get("parent") not in parent_set:
            continue
        code = post.get(opt_id)
        if code is None:
            continue
        items.append(
            OptimizeResponseItemSchema(
                source_code=code,
                explanation=explanations.get(opt_id, ""),
                optimization_id=opt_id,
                parent_id=meta.get("parent"),
            )
        )
    return items


def rerun_code_repair(record: OptimizationFeatures, parent_optimization_id: str) -> OptimizeResponseItemSchema | None:
    origin: dict[str, Any] = cast("dict[str, Any]", record.optimizations_origin) or {}
    post: dict[str, Any] = cast("dict[str, Any]", record.optimizations_post) or {}
    explanations: dict[str, Any] = cast("dict[str, Any]", record.explanations_post) or {}

    for opt_id, meta in origin.items():
        if meta.get("source") != "REPAIR":
            continue
        if meta.get("parent") != parent_optimization_id:
            continue
        code = post.get(opt_id)
        if code is None:
            continue
        return OptimizeResponseItemSchema(
            source_code=code,
            explanation=explanations.get(opt_id, ""),
            optimization_id=opt_id,
            parent_id=meta.get("parent"),
        )
    return None
