from typing import Dict, Any
from app.utils.llm_client import LLMClient

class BaseAgent:
    def __init__(self, name: str, template: str, client: LLMClient):
        self.name = name
        self.template = template
        self.client = client

    def render(self, **kwargs) -> str:
        return self.template.format(**kwargs)

    def invoke(self, **kwargs) -> Any:
        prompt = self.render(**kwargs)
        return self.client.invoke(prompt)

# --- Agent templates ---------------------------------------------------------

THREAT_MODELER_PROMPT = """You are a cybersecurity threat-modeling expert.
Input: {system_description}
1. Identify assets and trust boundaries.
2. Apply STRIDE categories.
3. Map to MITRE ATT&CK techniques.
Return JSON objects: asset, threat, stride_category, mitre_id, mitigation.
"""

THREAT_PATTERN_PROMPT = """Extract reusable threat patterns from this text.
Input: {text_chunk}
Return JSON list: pattern_name, description, preconditions, mitigations.
"""

ARCH_MAPPER_PROMPT = """Convert this architecture description into a DFD model.
Input: {components_and_interactions}
Return JSON: nodes, edges, trust_boundaries.
"""

# --- Agent classes -----------------------------------------------------------

class ThreatModelerAgent(BaseAgent):
    def run(self, system_description: str) -> Any:
        return self.invoke(system_description=system_description)

class ThreatPatternExtractor(BaseAgent):
    def run(self, text_chunk: str) -> Any:
        return self.invoke(text_chunk=text_chunk)

class ArchitectureMapper(BaseAgent):
    def run(self, components_and_interactions: str) -> Any:
        return self.invoke(components_and_interactions=components_and_interactions)

# --- Factory -----------------------------------------------------------------

def build_agents(model_name: str = "gpt-threat-intel-v1"):
    client = LLMClient(model_name)
    return {
        "modeler": ThreatModelerAgent("ThreatModelerAgent", THREAT_MODELER_PROMPT, client),
        "extractor": ThreatPatternExtractor("ThreatPatternExtractor", THREAT_PATTERN_PROMPT, client),
        "mapper": ArchitectureMapper("ArchitectureMapper", ARCH_MAPPER_PROMPT, client),
    }