import pytest

from pydantic_core import ValidationError

from aegis_ai.agents import rh_feature_agent
from aegis_ai.data_models import CVEID
from aegis_ai.features import component, cve
from tests.utils.llm_cache import llm_cache_retrieve

pytestmark = pytest.mark.asyncio


async def test_suggest_impact_with_test_model():
    def feature():
        return cve.SuggestImpact(rh_feature_agent).exec("CVE-2025-0725")

    result = await llm_cache_retrieve(feature)
    suggestimpact = cve.data_models.SuggestImpactModel.model_validate_json(result)
    assert isinstance(suggestimpact, cve.data_models.SuggestImpactModel)
    assert suggestimpact.impact == "LOW"
    assert (
        suggestimpact.disclaimer
        == "This response was generated by Aegis AI (https://github.com/RedHatProductSecurity/aegis-ai) using generative AI for informational purposes. All findings should be validated by a human expert."
    )


async def test_suggest_cwe_with_test_model():
    def feature():
        return cve.SuggestCWE(rh_feature_agent).exec("CVE-2025-0725")

    result = await llm_cache_retrieve(feature)
    suggestcwe = cve.data_models.SuggestCWEModel.model_validate_json(result)
    assert isinstance(suggestcwe, cve.data_models.SuggestCWEModel)
    assert suggestcwe.cwe == ["CWE-680"]


async def test_identify_pii_with_test_model():
    def feature():
        cve_id = CVEID(
            "CVE-2025-0725"
        )  # we can directly use custom fields though auto validation happens during feature input
        return cve.IdentifyPII(rh_feature_agent).exec(cve_id)

    result = await llm_cache_retrieve(feature)
    piireport = cve.data_models.PIIReportModel.model_validate_json(result)
    assert isinstance(piireport, cve.data_models.PIIReportModel)
    assert not piireport.contains_PII  # is false


async def test_rewrite_description_with_test_model():
    def feature():
        return cve.RewriteDescriptionText(rh_feature_agent).exec("CVE-2025-0725")

    result = await llm_cache_retrieve(feature)
    rewritedescription = cve.data_models.RewriteDescriptionModel.model_validate_json(
        result
    )
    assert isinstance(rewritedescription, cve.data_models.RewriteDescriptionModel)
    assert (
        rewritedescription.rewritten_title
        == "libcurl: gzip decompression buffer overflow"
    )


async def test_rewrite_statement_with_test_model():
    def feature():
        return cve.RewriteStatementText(rh_feature_agent).exec("CVE-2025-0725")

    result = await llm_cache_retrieve(feature)
    rewritestatement = cve.data_models.RewriteStatementModel.model_validate_json(result)
    assert isinstance(rewritestatement, cve.data_models.RewriteStatementModel)
    assert (
        rewritestatement.rewritten_statement
        == "This flaw is only exploitable when libcurl uses zlib version 1.2.0.3 or older for gzip decompression. Supported versions of Red Hat Enterprise Linux are not affected as they ship with newer versions of the zlib library. Some Red Hat services that bundle older third-party tools may be affected."
    )


async def test_cvss_diff_explain_with_test_model():
    def feature():
        return cve.CVSSDiffExplainer(rh_feature_agent).exec("CVE-2025-0725")

    result = await llm_cache_retrieve(feature)
    cvssdiffexplain = cve.data_models.CVSSDiffExplainerModel.model_validate_json(result)
    assert isinstance(cvssdiffexplain, cve.data_models.CVSSDiffExplainerModel)
    assert (
        cvssdiffexplain.redhat_cvss3_vector
        == "CVSS:3.1/AV:L/AC:L/PR:N/UI:N/S:U/C:N/I:N/A:L"
    )


async def test_component_intelligence_test_model():
    def feature():
        return component.ComponentIntelligence(rh_feature_agent).exec("curl")

    result = await llm_cache_retrieve(feature)
    componentintelligence = (
        component.data_models.ComponentIntelligenceModel.model_validate_json(result)
    )
    assert isinstance(
        componentintelligence, component.data_models.ComponentIntelligenceModel
    )
    assert componentintelligence.popularity_score == 1
    assert componentintelligence.confidence == 0.95


async def test_suggest_impact_with_bad_cve_test_model():
    def feature():
        return cve.SuggestImpact(rh_feature_agent).exec("BAD-CVE-ID")

    with pytest.raises(ValidationError) as excinfo:
        await feature()

    assert "String should match pattern" in str(excinfo)
