from abc import ABC, abstractmethod
from typing import Dict, Any, List
import asyncio
from pydantic import BaseModel, Field
from agents.llm_client import LLMClient

class AgentResponseSchema(BaseModel):
    agent: str = Field(description="Name of the agent, e.g., SecurityAgent")
    confidence: float = Field(description="Confidence score between 0.0 and 1.0 based on relevance and sufficiency of data")
    answer: str = Field(description="Detailed answer or analysis regarding the user query")
    citations: List[str] = Field(description="Source files, line numbers, or endpoints cited as reference")
    reasoning: List[str] = Field(description="Step-by-step reasoning steps the agent took")

class BaseAgent(ABC):
    def __init__(self, llm_client: LLMClient):
        self.llm_client = llm_client

    @abstractmethod
    async def run(
        self,
        profile: Dict[str, Any],
        graph: Dict[str, Any],
        summary: Dict[str, Any],
        report: str,
        query: str
    ) -> Dict[str, Any]:
        """Runs the agent's analysis."""
        pass

    async def _call_llm_json(self, prompt: str, schema: Any, temperature: float = 0.2) -> Dict[str, Any]:
        return await asyncio.to_thread(self.llm_client.generate_json, prompt, schema, temperature)