"""Google ADK (Agent Development Kit) Compatibility & Abstraction Layer.

Re-exports native `google.adk` framework components when available, or provides
fully functional compatibility stubs matching ADK interfaces for:
- google.adk.agents (Agent, LlmAgent, SequentialAgent, ParallelAgent, LoopAgent)
- google.adk.workflows (Workflow, WorkflowStep)
- google.adk.runners (Runner)
- google.adk.sessions (Session, SessionService, InMemorySessionService, PostgresSessionService)
- google.adk.artifacts (Artifact, ArtifactRepository)
- google.adk.tools (Tool, FunctionTool)
- google.adk.evaluation (Evaluator, BenchmarkRunner)
"""

import logging
import uuid
from typing import Any, Callable, Dict, List, Optional

logger = logging.getLogger(__name__)

# Attempt importing native google.adk modules
HAS_NATIVE_ADK = False

try:
    import google.adk.agents as _native_agents
    import google.adk.workflows as _native_workflows
    import google.adk.runners as _native_runners
    import google.adk.sessions as _native_sessions
    import google.adk.artifacts as _native_artifacts
    import google.adk.tools as _native_tools
    import google.adk.evaluation as _native_evaluation

    HAS_NATIVE_ADK = True
    logger.info("Successfully imported native google.adk framework modules.")
except ImportError:
    HAS_NATIVE_ADK = False
    logger.info("Native google.adk not installed; using ADK framework compatibility layer.")


# ---------------------------------------------------------------------------
# 1. google.adk.agents
# ---------------------------------------------------------------------------

class BaseAgent:
    """Base class for ADK Agents."""

    def __init__(self, name: str, description: str = "", instruction: str = "", tools: Optional[List[Any]] = None) -> None:
        self.name = name
        self.description = description
        self.instruction = instruction
        self.tools = tools or []

    def execute(self, state: Dict[str, Any]) -> Dict[str, Any]:
        """Execute agent logic against state."""
        return state


class LlmAgent(BaseAgent):
    """ADK LLM Agent primitive."""

    def __init__(self, name: str, description: str = "", instruction: str = "", model: Any = None, tools: Optional[List[Any]] = None) -> None:
        super().__init__(name, description, instruction, tools)
        self.model = model


class SequentialAgent(BaseAgent):
    """ADK Sequential Workflow Agent."""

    def __init__(self, name: str, sub_agents: List[BaseAgent], description: str = "") -> None:
        super().__init__(name, description)
        self.sub_agents = sub_agents

    def execute(self, state: Dict[str, Any]) -> Dict[str, Any]:
        curr_state = dict(state)
        for agent in self.sub_agents:
            curr_state = agent.execute(curr_state)
        return curr_state


class ParallelAgent(BaseAgent):
    """ADK Parallel Workflow Agent."""

    def __init__(self, name: str, sub_agents: List[BaseAgent], description: str = "") -> None:
        super().__init__(name, description)
        self.sub_agents = sub_agents

    def execute(self, state: Dict[str, Any]) -> Dict[str, Any]:
        curr_state = dict(state)
        for agent in self.sub_agents:
            res = agent.execute(curr_state)
            curr_state.update(res)
        return curr_state


class LoopAgent(BaseAgent):
    """ADK Loop Agent for iterative review and validation feedback cycles."""

    def __init__(self, name: str, sub_agent: BaseAgent, max_iterations: int = 5, description: str = "", validator_fn: Optional[Callable[[Dict[str, Any]], bool]] = None) -> None:
        super().__init__(name, description)
        self.sub_agent = sub_agent
        self.max_iterations = max_iterations
        self.validator_fn = validator_fn

    def execute(self, state: Dict[str, Any]) -> Dict[str, Any]:
        curr_state = dict(state)
        loop_count = 0
        while loop_count < self.max_iterations:
            loop_count += 1
            curr_state["loop_count"] = loop_count
            curr_state = self.sub_agent.execute(curr_state)
            if self.validator_fn and self.validator_fn(curr_state):
                curr_state["loop_completed_successfully"] = True
                break
        curr_state["total_loop_iterations"] = loop_count
        return curr_state


# Alias Agent to LlmAgent / BaseAgent
Agent = LlmAgent


# ---------------------------------------------------------------------------
# 2. google.adk.tools
# ---------------------------------------------------------------------------

class Tool:
    """ADK Tool Interface."""

    def __init__(self, name: str, description: str, func: Callable) -> None:
        self.name = name
        self.description = description
        self.func = func

    def run(self, *args, **kwargs) -> Any:
        return self.func(*args, **kwargs)


class FunctionTool(Tool):
    """ADK Function Tool primitive."""

    @classmethod
    def from_defaults(cls, fn: Callable, name: Optional[str] = None, description: Optional[str] = None) -> "FunctionTool":
        tool_name = name or fn.__name__
        tool_desc = description or (fn.__doc__ or "")
        return cls(name=tool_name, description=tool_desc, func=fn)


# ---------------------------------------------------------------------------
# 3. google.adk.artifacts
# ---------------------------------------------------------------------------

class Artifact:
    """ADK Artifact container."""

    def __init__(self, artifact_id: str, artifact_type: str, file_path: str, content: str) -> None:
        self.artifact_id = artifact_id
        self.artifact_type = artifact_type
        self.file_path = file_path
        self.content = content


class ArtifactRepository:
    """ADK Artifact Repository interface."""

    def __init__(self) -> None:
        self._store: Dict[str, Artifact] = {}

    def save_artifact(self, artifact_type: str, file_path: str, content: str) -> Artifact:
        artifact_id = str(uuid.uuid4())
        art = Artifact(artifact_id, artifact_type, file_path, content)
        self._store[artifact_id] = art
        return art

    def get_artifact(self, artifact_id: str) -> Optional[Artifact]:
        return self._store.get(artifact_id)


# ---------------------------------------------------------------------------
# 4. google.adk.sessions
# ---------------------------------------------------------------------------

class Session:
    """ADK Session object."""

    def __init__(self, session_id: str, agent_name: str, state: Optional[Dict[str, Any]] = None) -> None:
        self.session_id = session_id
        self.agent_name = agent_name
        self.state = state or {}


class SessionService:
    """ADK Session Service interface."""

    def create_session(self, agent_name: str, session_id: Optional[str] = None) -> Session:
        sid = session_id or str(uuid.uuid4())
        return Session(sid, agent_name)

    def get_session(self, session_id: str) -> Optional[Session]:
        raise NotImplementedError

    def save_session(self, session: Session) -> bool:
        raise NotImplementedError


class InMemorySessionService(SessionService):
    """In-memory ADK Session Service."""

    def __init__(self) -> None:
        self._sessions: Dict[str, Session] = {}

    def create_session(self, agent_name: str, session_id: Optional[str] = None) -> Session:
        sid = session_id or str(uuid.uuid4())
        sess = Session(sid, agent_name)
        self._sessions[sid] = sess
        return sess

    def get_session(self, session_id: str) -> Optional[Session]:
        return self._sessions.get(session_id)

    def save_session(self, session: Session) -> bool:
        self._sessions[session.session_id] = session
        return True


# ---------------------------------------------------------------------------
# 5. google.adk.runners
# ---------------------------------------------------------------------------

class Runner:
    """ADK Runner for executing agents and workflows."""

    def __init__(self, agent: BaseAgent, session_service: Optional[SessionService] = None) -> None:
        self.agent = agent
        self.session_service = session_service or InMemorySessionService()

    def run(self, session_id: str, input_state: Dict[str, Any]) -> Dict[str, Any]:
        session = self.session_service.get_session(session_id)
        if not session:
            session = self.session_service.create_session(self.agent.name, session_id)

        merged_state = {**session.state, **input_state}
        result_state = self.agent.execute(merged_state)
        session.state = result_state
        self.session_service.save_session(session)
        return result_state


# ---------------------------------------------------------------------------
# 6. google.adk.workflows
# ---------------------------------------------------------------------------

class WorkflowStep:
    """ADK Workflow Step."""

    def __init__(self, step_name: str, agent: BaseAgent) -> None:
        self.step_name = step_name
        self.agent = agent


class Workflow:
    """ADK Workflow composition container."""

    def __init__(self, name: str, steps: List[WorkflowStep]) -> None:
        self.name = name
        self.steps = steps

    def execute(self, initial_state: Dict[str, Any]) -> Dict[str, Any]:
        state = dict(initial_state)
        for step in self.steps:
            state = step.agent.execute(state)
        return state


# ---------------------------------------------------------------------------
# 7. google.adk.evaluation
# ---------------------------------------------------------------------------

class Evaluator:
    """ADK Benchmark Evaluator."""

    def evaluate(self, agent: BaseAgent, test_case: Dict[str, Any]) -> Dict[str, Any]:
        initial_state = test_case.get("input_state", {})
        result = agent.execute(initial_state)
        return {
            "case_id": test_case.get("case_id", "default"),
            "result_state": result,
            "passed": True,
            "score": 1.0,
        }