from enum import Enum

from pydantic import BaseModel, Field

from broadie import ToolResponse, create_agent, create_sub_agent, tool

# ---------------------------
# Custom tools for this agent
# ---------------------------


@tool("lookup_ip", description="Look up IP reputation in a threat intel database")
def lookup_ip(ip: str) -> ToolResponse:
    # fake lookup
    reputation = "malicious" if ip.startswith("192.") else "clean"
    return ToolResponse.success(
        message=f"IP {ip} reputation lookup completed",
        data={"ip": ip, "reputation": reputation},
        meta={"source": "threat_intel_db", "lookup_type": "ip"},
    )


@tool("lookup_domain", description="Check domain reputation in DNS and threat feeds")
def lookup_domain(domain: str) -> ToolResponse:
    category = "phishing" if "phish" in domain else "benign"
    return ToolResponse.success(
        message=f"Domain {domain} reputation lookup completed",
        data={"domain": domain, "category": category},
        meta={"source": "dns_threat_feeds", "lookup_type": "domain"},
    )


# ---------------------------
# Output schemas
# ---------------------------


class Verdict(str, Enum):
    malicious = "malicious"
    suspicious = "suspicious"
    benign = "benign"
    unknown = "unknown"


class EnrichmentResult(BaseModel):
    indicator: str
    type: str
    verdict: Verdict


class ThreatIntelOutput(BaseModel):
    summary: str = Field(
        ...,
        description="Human-readable summary of the threat analysis",
    )
    enrichments: list[EnrichmentResult] = Field(
        ...,
        description="List of enriched indicators with verdicts",
    )


# ---------------------------
# SubAgent
# ---------------------------

indicator_enricher = create_sub_agent(
    name="indicator_enricher",
    prompt="Given an IP or domain, enrich it with threat intel context.",
    description="Enriches security indicators (IPs, domains) with threat intelligence data",
    output_schema=EnrichmentResult,
    tools=[lookup_ip, lookup_domain],
)

# ---------------------------
# Main Agent
# ---------------------------

threat_intel_agent = create_agent(
    name="threat_intel",
    instruction="You are a threat intelligence assistant. Analyze indicators and produce structured enrichment.",
    output_schema=ThreatIntelOutput,
    tools=[lookup_ip, lookup_domain],
    subagents=[indicator_enricher],
    channels=[
        {
            "type": "slack",
            "target": "#threat-feed",
            "instructions": "Summarize in Slack blocks",
        },
    ],
)


if __name__ == "__main__":
    import asyncio

    async def main():
        email_message = "Hello user, please reset your password here: http://bad-link.com"
        result = await threat_intel_agent.run(email_message)
        print(result)

    asyncio.run(main())
