#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
# See LICENSE file for details.
"""
PULSE Worker LLM — client initialization, API calls, and tool result handling.

Supports: Anthropic, OpenAI, Google Gemini, Ollama.
Imported by pulse_worker.py.
"""
from __future__ import annotations

import json
import os
from pathlib import Path

from .worker_tools import TOOL_DEFINITIONS_ANTHROPIC, TOOL_DISPATCH


class WorkerLLM:
    """Manages LLM client lifecycle, API calls, and tool result formatting."""

    def __init__(self, provider: str, model: str):
        self.provider = provider
        self.model = model
        self._client = None

    def init_client(self):
        """Lazy-init the LLM client."""
        if self._client is not None:
            return

        if self.provider == "anthropic":
            import anthropic
            self._client = anthropic.Anthropic()
        elif self.provider == "google":
            from dotenv import load_dotenv
            load_dotenv()
            api_key = os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY") or ""
            try:
                from google import genai
                self._client = genai.Client(api_key=api_key)
                self._google_sdk = "new"
            except ImportError:
                import google.generativeai as genai
                genai.configure(api_key=api_key)
                self._client = genai
                self._google_sdk = "old"
        elif self.provider == "openai":
            import openai
            self._client = openai.OpenAI()
        elif self.provider == "ollama":
            self._client = None  # HTTP-based, no SDK needed
            self._ollama_endpoint = os.environ.get(
                "OLLAMA_ENDPOINT", "http://localhost:11434"
            )
        else:
            raise ValueError(f"Unknown provider: {self.provider}")

    # ------------------------------------------------------------------
    # Provider-specific API calls
    # ------------------------------------------------------------------

    def _call_anthropic(self, system: str, messages: list) -> tuple[str, list, dict]:
        resp = self._client.messages.create(
            model=self.model, system=system, messages=messages,
            tools=TOOL_DEFINITIONS_ANTHROPIC, max_tokens=4096,
        )
        text_parts = []
        tool_calls = []
        for block in resp.content:
            if block.type == "text":
                text_parts.append(block.text)
            elif block.type == "tool_use":
                tool_calls.append({"id": block.id, "name": block.name, "input": block.input})
        return "\n".join(text_parts), tool_calls, {"in": resp.usage.input_tokens, "out": resp.usage.output_tokens}

    def _call_openai(self, system: str, messages: list) -> tuple[str, list, dict]:
        oai_messages = [{"role": "system", "content": system}]
        for m in messages:
            if m["role"] == "tool":
                oai_messages.append({"role": "tool", "tool_call_id": m.get("tool_call_id", ""), "content": m["content"]})
            elif m["role"] == "assistant" and "tool_calls" in m:
                oai_messages.append(m)
            else:
                oai_messages.append({"role": m["role"], "content": m.get("content", "")})

        oai_tools = [
            {"type": "function", "function": {"name": t["name"], "description": t["description"], "parameters": t["input_schema"]}}
            for t in TOOL_DEFINITIONS_ANTHROPIC
        ]
        resp = self._client.chat.completions.create(
            model=self.model, messages=oai_messages, tools=oai_tools, max_tokens=4096,
        )
        choice = resp.choices[0]
        text = choice.message.content or ""
        tool_calls = []
        if choice.message.tool_calls:
            for tc in choice.message.tool_calls:
                tool_calls.append({"id": tc.id, "name": tc.function.name, "input": json.loads(tc.function.arguments)})
        return text, tool_calls, {"in": resp.usage.prompt_tokens, "out": resp.usage.completion_tokens}

    def _call_google(self, system: str, messages: list) -> tuple[str, list, dict]:
        from google.genai import types

        func_declarations = []
        for t in TOOL_DEFINITIONS_ANTHROPIC:
            props = {}
            required = t["input_schema"].get("required", [])
            for pname, pdef in t["input_schema"].get("properties", {}).items():
                props[pname] = types.Schema(type=types.Type.STRING, description=pdef.get("description", ""))
            func_declarations.append(types.FunctionDeclaration(
                name=t["name"], description=t["description"],
                parameters=types.Schema(type=types.Type.OBJECT, properties=props, required=required),
            ))
        tool = types.Tool(function_declarations=func_declarations)

        # Build Google-compatible content history.
        # Key rule: function_call parts go in model turn, function_response
        # parts go in the IMMEDIATELY FOLLOWING user turn. Never separate them.
        contents = []
        for m in messages:
            if m["role"] == "user":
                contents.append(types.Content(role="user", parts=[types.Part.from_text(text=m["content"])]))
            elif m["role"] == "model_with_calls":
                # Combined model turn: text + function calls (built by append_tool_results)
                parts = []
                if m.get("text"):
                    parts.append(types.Part.from_text(text=m["text"]))
                for fc in m.get("function_calls", []):
                    parts.append(types.Part.from_function_call(name=fc["name"], args=fc["args"]))
                if parts:
                    contents.append(types.Content(role="model", parts=parts))
            elif m["role"] == "function_responses":
                # Batch all function responses into one user turn
                parts = [
                    types.Part.from_function_response(name=r["name"], response={"result": r["content"]})
                    for r in m["responses"]
                ]
                contents.append(types.Content(role="user", parts=parts))
            elif m["role"] == "assistant":
                # Plain text model response (no tool calls)
                contents.append(types.Content(role="model", parts=[types.Part.from_text(text=m.get("content", "") or ".")]))

        resp = self._client.models.generate_content(
            model=self.model, contents=contents,
            config=types.GenerateContentConfig(system_instruction=system, tools=[tool]),
        )
        text_parts = []
        tool_calls = []
        if resp.candidates and resp.candidates[0].content:
            for part in resp.candidates[0].content.parts:
                if part.text:
                    text_parts.append(part.text)
                elif part.function_call:
                    fc = part.function_call
                    tool_calls.append({"id": f"gc_{fc.name}_{len(tool_calls)}", "name": fc.name, "input": dict(fc.args) if fc.args else {}})
        usage = {"in": 0, "out": 0}
        if hasattr(resp, "usage_metadata"):
            usage["in"] = resp.usage_metadata.prompt_token_count
            usage["out"] = resp.usage_metadata.candidates_token_count
        return "\n".join(text_parts), tool_calls, usage

    def _call_ollama(self, system: str, messages: list) -> tuple[str, list, dict]:
        import urllib.request
        import urllib.error

        ollama_messages = [{"role": "system", "content": system}]
        for m in messages:
            if m["role"] == "tool_result":
                ollama_messages.append({"role": "tool", "content": m.get("content", "")})
            elif m["role"] == "assistant" and isinstance(m.get("content"), list):
                text_parts = []
                for block in m["content"]:
                    if isinstance(block, dict) and block.get("type") == "text":
                        text_parts.append(block["text"])
                ollama_messages.append({"role": "assistant", "content": "\n".join(text_parts)})
            else:
                ollama_messages.append({
                    "role": m["role"],
                    "content": m.get("content", "") if isinstance(m.get("content"), str) else str(m.get("content", "")),
                })

        ollama_tools = [
            {"type": "function", "function": {"name": t["name"], "description": t["description"], "parameters": t["input_schema"]}}
            for t in TOOL_DEFINITIONS_ANTHROPIC
        ]

        payload = {"model": self.model, "messages": ollama_messages, "tools": ollama_tools, "stream": False}
        url = f"{self._ollama_endpoint}/api/chat"
        data = json.dumps(payload).encode("utf-8")
        req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"}, method="POST")

        try:
            with urllib.request.urlopen(req, timeout=300) as resp:
                result = json.loads(resp.read().decode("utf-8"))
        except urllib.error.URLError as e:
            raise ConnectionError(f"Ollama not reachable at {self._ollama_endpoint}: {e}")

        msg = result.get("message", {})
        text = msg.get("content", "") or ""
        tool_calls = []
        for tc in msg.get("tool_calls", []):
            func = tc.get("function", {})
            tool_calls.append({
                "id": f"ollama_{func.get('name', '')}_{len(tool_calls)}",
                "name": func.get("name", ""),
                "input": func.get("arguments", {}),
            })
        usage = {"in": result.get("prompt_eval_count", 0), "out": result.get("eval_count", 0)}
        return text, tool_calls, usage

    # ------------------------------------------------------------------
    # Unified call + tool execution
    # ------------------------------------------------------------------

    def call_llm(self, system: str, messages: list) -> tuple[str, list, dict]:
        """Route to the correct provider."""
        if self.provider == "anthropic":
            return self._call_anthropic(system, messages)
        elif self.provider == "openai":
            return self._call_openai(system, messages)
        elif self.provider == "google":
            return self._call_google(system, messages)
        elif self.provider == "ollama":
            return self._call_ollama(system, messages)
        raise ValueError(f"Unknown provider: {self.provider}")

    def run_tools(self, tool_calls: list) -> list[dict]:
        """Execute tool calls and return results."""
        import logging
        log = logging.getLogger("pulse.worker_llm")
        results = []
        for tc in tool_calls:
            name = tc["name"]
            args = tc["input"]
            handler = TOOL_DISPATCH.get(name)
            try:
                output = handler(args) if handler else f"Unknown tool: {name}"
            except Exception as e:
                output = f"TOOL ERROR: {e}"
                log.warning("Tool %s raised: %s", name, e)
            raw = str(output)
            if len(raw) > 10000:
                log.debug("Tool %s output truncated: %d -> 10000 chars", name, len(raw))
            results.append({"id": tc["id"], "name": name, "content": raw[:10000]})
        return results

    def append_tool_results(self, messages: list, text: str, tool_calls: list, results: list):
        """Append assistant + tool results to message history in provider-specific format."""
        if self.provider == "anthropic":
            assistant_content = []
            if text:
                assistant_content.append({"type": "text", "text": text})
            for tc in tool_calls:
                assistant_content.append({"type": "tool_use", "id": tc["id"], "name": tc["name"], "input": tc["input"]})
            messages.append({"role": "assistant", "content": assistant_content})
            messages.append({"role": "user", "content": [{"type": "tool_result", "tool_use_id": r["id"], "content": r["content"]} for r in results]})
        elif self.provider == "openai":
            messages.append({
                "role": "assistant", "content": text or None,
                "tool_calls": [{"id": tc["id"], "type": "function", "function": {"name": tc["name"], "arguments": json.dumps(tc["input"])}} for tc in tool_calls],
            })
            for r in results:
                messages.append({"role": "tool", "tool_call_id": r["id"], "content": r["content"]})
        elif self.provider == "google":
            # Google requires: model turn with function_call parts, then
            # immediately a user turn with function_response parts.
            # Both must be single turns (not split across messages).
            messages.append({
                "role": "model_with_calls",
                "text": text or "",
                "function_calls": [{"name": tc["name"], "args": tc["input"]} for tc in tool_calls],
            })
            messages.append({
                "role": "function_responses",
                "responses": [{"name": r["name"], "content": r["content"]} for r in results],
            })
        elif self.provider == "ollama":
            messages.append({
                "role": "assistant", "content": text or "",
                "tool_calls": [{"id": tc["id"], "type": "function", "function": {"name": tc["name"], "arguments": json.dumps(tc["input"])}} for tc in tool_calls],
            })
            for r in results:
                messages.append({"role": "tool", "content": r["content"]})
