import json
import os
import re
import subprocess
import uuid

from hiveos.agent_interface import AgentInterface
from hiveos.config import AI_OLLAMA_MODEL, AI_OLLAMA_TIMEOUT_SEC, AI_PROVIDER
from hiveos.intelligence.contracts import validate_ai_result, validate_tool_request
from hiveos.intelligence.tracing import build_trace_event, emit_trace, now_iso


class AIAgent(AgentInterface):
    def __init__(self, agent_id=None, name=None, capabilities=None):
        super().__init__()
        self.agent_id = agent_id or f"ai-agent-{uuid.uuid4().hex[:8]}"
        self.name = name or f"AIAgent-{self.agent_id[-4:]}"
        self.capabilities = capabilities or ["ai"]
        self.latest_ai_result = None
        self.last_error = None
        self.current_run_id = None
        self.execution_plane = None
        self.tool_registry = None
        self.proposal_manager = None
        self.proposal_core = None
        self.latest_tool_result = None
        self.latest_proposal = None

    def get_status(self):
        return self.status

    def report_capabilities(self):
        return self.capabilities

    def execute_chunk(self, chunk, payload=None):
        super().execute_chunk(chunk, payload)
        if self.execution_plane is not None:
            self.current_run_id = self.execution_plane.submit_run(
                agent_id=self.agent_id,
                task_id=chunk.task_id,
                chunk_id=chunk.chunk_id,
                payload=self.payload,
            )
        else:
            self.current_run_id = f"{chunk.chunk_id}:{uuid.uuid4().hex[:8]}"
        normalized = self._normalize_payload_to_ai_result(chunk, self.payload)
        self.latest_ai_result = validate_ai_result(normalized)
        self.publish(f"ai/{self.agent_id}/result", self.latest_ai_result)
        emit_trace(
            build_trace_event(
                event_type="ai_output",
                run_id=self.current_run_id,
                outcome="success",
                payload={"type": self.latest_ai_result["type"]},
                agent_id=self.agent_id,
                task_id=chunk.task_id,
                chunk_id=chunk.chunk_id,
                zone_id=self.zone,
                model=self.latest_ai_result.get("meta", {}),
            )
        )
        if self.latest_ai_result["type"] == "tool_request":
            self._handle_tool_request(chunk)
        elif self.latest_ai_result["type"] == "proposal":
            self._handle_proposal(chunk)

    def _normalize_payload_to_ai_result(self, chunk, payload):
        if isinstance(payload, dict) and payload.get("type") in {
            "reply",
            "observation",
            "proposal",
            "tool_request",
            "error",
        }:
            candidate = dict(payload)
        elif isinstance(payload, dict) and self._is_proposal_prompt_payload(payload):
            candidate = self._build_proposal_result_from_prompt(chunk, payload)
        else:
            text = payload if isinstance(payload, str) else str(payload)
            candidate = {
                "type": "reply",
                "content": {"text": text, "audience": "system"},
            }

        candidate.setdefault("agent_id", self.agent_id)
        candidate.setdefault("task_id", chunk.task_id)
        candidate.setdefault("chunk_id", chunk.chunk_id)
        candidate.setdefault("run_id", self.current_run_id or f"{chunk.chunk_id}:init")
        candidate.setdefault("timestamp", now_iso())
        candidate.setdefault("meta", {"provider": "local", "model": "stub"})
        return candidate

    def _is_proposal_prompt_payload(self, payload):
        mode = str(payload.get("ai_mode", "")).lower()
        intent = str(payload.get("intent", "")).lower()
        return (
            mode == "proposal"
            or intent in {"proposal", "generate_proposal", "proposal_draft"}
            or "compose_prompt" in payload
            or "proposal_prompt" in payload
        )

    def _build_proposal_result_from_prompt(self, chunk, payload):
        prompt_text = (
            payload.get("compose_prompt")
            or payload.get("proposal_prompt")
            or payload.get("prompt")
            or ""
        )
        prompt_text = str(prompt_text).strip() or f"Generate proposal for chunk {chunk.chunk_id}"

        model = str(AI_OLLAMA_MODEL or "phi3")
        timeout_sec = _safe_int(str(AI_OLLAMA_TIMEOUT_SEC or 40), 40)
        provider = str(AI_PROVIDER or "ollama").lower()

        if provider == "ollama":
            raw_output = self._run_ollama_proposal(prompt_text, model, timeout_sec)
            artifact = self._extract_or_fallback_artifact(raw_output, prompt_text, chunk, payload)
            return {
                "type": "proposal",
                "content": artifact,
                "meta": {"provider": "ollama", "model": model},
            }

        # Unknown provider fallback still creates a valid proposal artifact for testability.
        artifact = self._fallback_artifact(prompt_text, chunk, payload, summary="unknown_provider_fallback")
        return {
            "type": "proposal",
            "content": artifact,
            "meta": {"provider": provider, "model": "fallback"},
        }

    def _run_ollama_proposal(self, prompt_text, model, timeout_sec):
        instructions = (
            "You are a HiveOS proposal generator. "
            "Return ONLY JSON object with keys: artifact_type, title, summary, body. "
            "artifact_type must be task_draft. "
            "body must include: parent_task_id (int), depth (int), auto_inject (bool), "
            "task_intent { task_name, execution_mode, capability_required, chunks[] }. "
            "depth must be 1. auto_inject must be false. "
            "execution_mode must be sequential or concurrent. "
            "capability_required must be a list of strings. "
            "chunks must be a list of objects, each with chunk_id, required_capability, payload, ttl, timing. "
            "Keep exactly one chunk in task_intent.chunks for now."
        )
        full_prompt = f"{instructions}\n\nUSER_GOAL:\n{prompt_text}"
        try:
            result = subprocess.run(
                ["ollama", "run", model, full_prompt],
                capture_output=True,
                text=True,
                encoding="utf-8",
                timeout=max(5, timeout_sec),
            )
            return (result.stdout or "").strip()
        except Exception as exc:
            self.last_error = f"ollama_error:{exc}"
            return ""

    def _extract_or_fallback_artifact(self, raw_output, prompt_text, chunk, payload):
        parsed = _extract_json_object(raw_output)
        if isinstance(parsed, dict):
            return self._normalize_artifact(parsed, prompt_text, chunk, payload)
        return self._fallback_artifact(prompt_text, chunk, payload, summary="parse_fallback")

    def _normalize_artifact(self, artifact, prompt_text, chunk, payload):
        body = artifact.get("body")
        if not isinstance(body, dict):
            body = {}
        task_intent = _sanitize_task_intent(
            body.get("task_intent"),
            prompt_text=prompt_text,
            chunk=chunk,
            payload=payload,
        )

        depth = _safe_int(body.get("depth"), 1)
        if depth < 1:
            depth = 1
        if depth > 1:
            depth = 1

        parent_task_id = _safe_int(body.get("parent_task_id"), int(chunk.task_id))

        return {
            "artifact_type": "task_draft",
            "title": str(artifact.get("title") or f"AI proposal for {chunk.chunk_id}"),
            "summary": str(artifact.get("summary") or "Generated from local model."),
            "body": {
                "parent_task_id": parent_task_id,
                "depth": depth,
                "auto_inject": False,
                "task_intent": task_intent,
            },
        }

    def _fallback_artifact(self, prompt_text, chunk, payload, summary="fallback"):
        return {
            "artifact_type": "task_draft",
            "title": f"Fallback proposal for {chunk.chunk_id}",
            "summary": summary,
            "body": {
                "parent_task_id": int(chunk.task_id),
                "depth": 1,
                "auto_inject": False,
                "task_intent": _default_task_intent(prompt_text, chunk, payload),
            },
        }

    def report_completion(self, chunk):
        super().report_completion(chunk)

    def bind_execution_plane(self, execution_plane):
        self.execution_plane = execution_plane

    def bind_tool_registry(self, tool_registry):
        self.tool_registry = tool_registry

    def bind_proposal_manager(self, proposal_manager, core):
        self.proposal_manager = proposal_manager
        self.proposal_core = core

    def _handle_tool_request(self, chunk):
        content = self.latest_ai_result.get("content", {})
        try:
            request = {
                "request_id": f"tool-{uuid.uuid4().hex[:10]}",
                "run_id": self.current_run_id or f"{chunk.chunk_id}:tool",
                "agent_id": self.agent_id,
                "task_id": chunk.task_id,
                "chunk_id": chunk.chunk_id,
                "tool_id": content.get("tool_id"),
                "action": content.get("action"),
                "args": content.get("args", {}),
                "requested_scope": content.get("scope"),
                "budget_cost": content.get("budget_cost", 0),
                "timestamp": now_iso(),
            }
            request = validate_tool_request(request)
        except Exception as exc:
            self.last_error = f"invalid_tool_request:{exc}"
            emit_trace(
                build_trace_event(
                    event_type="tool_decision",
                    run_id=self.current_run_id or f"{chunk.chunk_id}:tool",
                    outcome="failed",
                    payload={"reason": self.last_error},
                    agent_id=self.agent_id,
                    task_id=chunk.task_id,
                    chunk_id=chunk.chunk_id,
                    zone_id=self.zone,
                )
            )
            return

        if self.tool_registry is None:
            self.last_error = "tool_registry_not_bound"
            emit_trace(
                build_trace_event(
                    event_type="tool_decision",
                    run_id=request["run_id"],
                    outcome="denied",
                    payload={"reason": self.last_error, "tool_id": request["tool_id"]},
                    agent_id=self.agent_id,
                    task_id=chunk.task_id,
                    chunk_id=chunk.chunk_id,
                    zone_id=self.zone,
                )
            )
            return

        decision = self.tool_registry.authorize(self, request)
        decision_outcome = "success" if decision.get("allowed") else "denied"
        emit_trace(
            build_trace_event(
                event_type="tool_decision",
                run_id=request["run_id"],
                outcome=decision_outcome,
                payload={
                    "tool_id": request["tool_id"],
                    "action": request["action"],
                    "reason": decision.get("reason"),
                },
                agent_id=self.agent_id,
                task_id=chunk.task_id,
                chunk_id=chunk.chunk_id,
                zone_id=self.zone,
            )
        )

        if not decision.get("allowed"):
            self.last_error = decision.get("reason")
            self.latest_tool_result = {"status": "denied", "reason": self.last_error}
            return

        result = self.tool_registry.execute(
            request,
            context={"chunk": chunk, "agent": self, "run_id": request["run_id"]},
        )
        self.latest_tool_result = {"status": "ok", "result": result}
        self.publish(f"ai/{self.agent_id}/tool_result", self.latest_tool_result)
        emit_trace(
            build_trace_event(
                event_type="tool_result",
                run_id=request["run_id"],
                outcome="success",
                payload={"tool_id": request["tool_id"], "result": result},
                agent_id=self.agent_id,
                task_id=chunk.task_id,
                chunk_id=chunk.chunk_id,
                zone_id=self.zone,
            )
        )

    def _handle_proposal(self, chunk):
        content = self.latest_ai_result.get("content", {})
        if self.proposal_manager is None:
            self.last_error = "proposal_manager_not_bound"
            emit_trace(
                build_trace_event(
                    event_type="proposal_created",
                    run_id=self.current_run_id or f"{chunk.chunk_id}:proposal",
                    outcome="failed",
                    payload={"reason": self.last_error},
                    agent_id=self.agent_id,
                    task_id=chunk.task_id,
                    chunk_id=chunk.chunk_id,
                    zone_id=self.zone,
                )
            )
            return

        created = self.proposal_manager.create(
            agent_id=self.agent_id,
            run_id=self.current_run_id or f"{chunk.chunk_id}:proposal",
            task_id=chunk.task_id,
            chunk_id=chunk.chunk_id,
            proposal_content=content,
        )
        self.latest_proposal = created
        emit_trace(
            build_trace_event(
                event_type="proposal_created",
                run_id=created["run_id"],
                outcome="success",
                payload={"proposal_id": created["proposal_id"], "artifact_type": content.get("artifact_type")},
                agent_id=self.agent_id,
                task_id=chunk.task_id,
                chunk_id=chunk.chunk_id,
                zone_id=self.zone,
            )
        )

        validation = self.proposal_manager.validate(
            created["proposal_id"],
            validation_mode="core_auto",
            validation_source="chain",
            validation_actor="system",
        )
        outcome = "success" if validation.get("ok") else "denied"
        emit_trace(
            build_trace_event(
                event_type="proposal_validated",
                run_id=created["run_id"],
                outcome=outcome,
                payload={"proposal_id": created["proposal_id"], "validation": validation},
                agent_id=self.agent_id,
                task_id=chunk.task_id,
                chunk_id=chunk.chunk_id,
                zone_id=self.zone,
            )
        )

        auto_inject = bool((content.get("body") or {}).get("auto_inject", False))
        if auto_inject and validation.get("ok") and self.proposal_core is not None:
            injected = self.proposal_manager.inject_validated(created["proposal_id"], self.proposal_core)
            self.latest_proposal = injected
            emit_trace(
                build_trace_event(
                    event_type="proposal_injected",
                    run_id=created["run_id"],
                    outcome="success",
                    payload={
                        "proposal_id": created["proposal_id"],
                        "child_task_id": injected.get("child_task_id"),
                        "parent_task_id": injected.get("parent_task_id"),
                    },
                    agent_id=self.agent_id,
                    task_id=chunk.task_id,
                    chunk_id=chunk.chunk_id,
                    zone_id=self.zone,
                )
            )

    def health_check(self):
        return {"status": self.status, "last_error": self.last_error}

    def enter_cooldown(self, reason: str = ""):
        self.status = "cooldown"
        self._available = False
        self.last_error = reason

    def is_available(self):
        return self._available and self.status == "idle"

    def get_summary(self):
        return super().get_summary() | {
            "type": "AIAgent",
            "last_result_type": (self.latest_ai_result or {}).get("type"),
            "current_run_id": self.current_run_id,
            "last_error": self.last_error,
            "latest_tool_result": self.latest_tool_result,
            "latest_proposal_id": (self.latest_proposal or {}).get("proposal_id"),
        }


def _safe_int(raw, default_value):
    try:
        return int(raw)
    except Exception:
        return default_value


def _extract_json_object(text):
    source = str(text or "").strip()
    if not source:
        return None

    fenced = re.search(r"```json\s*([\s\S]+?)\s*```", source, re.IGNORECASE)
    if fenced:
        source = fenced.group(1).strip()

    try:
        return json.loads(source)
    except Exception:
        pass

    first = source.find("{")
    last = source.rfind("}")
    if first >= 0 and last > first:
        try:
            return json.loads(source[first:last + 1])
        except Exception:
            return None
    return None


def _default_task_intent(prompt_text, chunk, payload):
    capability = str(payload.get("preferred_capability") or "compute.generic")
    return {
        "task_name": f"proposal_child_{chunk.chunk_id}",
        "execution_mode": "sequential",
        "capability_required": [capability],
        "chunks": [
            {
                "chunk_id": "proposal_step_1",
                "required_capability": capability,
                "payload": {
                    "intent": "execute_proposed_flow",
                    "parameters": {"prompt": prompt_text},
                },
                "ttl": 20,
                "timing": "ticks",
            }
        ],
    }


def _sanitize_task_intent(candidate, prompt_text, chunk, payload):
    default = _default_task_intent(prompt_text, chunk, payload)
    if not isinstance(candidate, dict):
        return default

    task_name = str(candidate.get("task_name") or default["task_name"])
    execution_mode = str(candidate.get("execution_mode") or default["execution_mode"]).lower()
    if execution_mode not in {"sequential", "concurrent"}:
        execution_mode = default["execution_mode"]

    required_caps = candidate.get("capability_required")
    if isinstance(required_caps, str):
        required_caps = [required_caps]
    if not isinstance(required_caps, list):
        required_caps = list(default["capability_required"])
    required_caps = [str(cap) for cap in required_caps if str(cap).strip()]
    if not required_caps:
        required_caps = list(default["capability_required"])

    chunks = candidate.get("chunks")
    if not isinstance(chunks, list):
        chunks = []
    sanitized_chunks = []
    for idx, raw in enumerate(chunks):
        if not isinstance(raw, dict):
            continue
        chunk_id = str(raw.get("chunk_id") or f"proposal_step_{idx + 1}")
        required_capability = str(raw.get("required_capability") or required_caps[0])
        payload_obj = raw.get("payload")
        if not isinstance(payload_obj, dict):
            payload_obj = {
                "intent": "execute_proposed_flow",
                "parameters": {"prompt": prompt_text},
            }
        ttl = _safe_int(raw.get("ttl"), 20)
        if ttl < 1:
            ttl = 1
        timing = str(raw.get("timing") or "ticks")
        if timing not in {"ticks", "seconds", "ms"}:
            timing = "ticks"
        sanitized_chunks.append(
            {
                "chunk_id": chunk_id,
                "required_capability": required_capability,
                "payload": payload_obj,
                "ttl": ttl,
                "timing": timing,
            }
        )
        if len(sanitized_chunks) >= 3:
            break

    if not sanitized_chunks:
        sanitized_chunks = list(default["chunks"])

    return {
        "task_name": task_name,
        "execution_mode": execution_mode,
        "capability_required": required_caps,
        "chunks": sanitized_chunks,
    }
