from __future__ import annotations

import os
from pathlib import Path
from typing import Any, Callable

from autogen_core.models import CreateResult, SystemMessage, UserMessage

from prompt_auto_lab.core.model_clients import build_papyrus_model_client


def _split_think_blocks(text: str) -> tuple[str, str | None]:
    """Return (visible_text, concatenated_thought) for <think>...</think> blocks."""
    visible_parts: list[str] = []
    thought_parts: list[str] = []
    i = 0
    while True:
        start = text.find("<think>", i)
        if start == -1:
            visible_parts.append(text[i:])
            break
        visible_parts.append(text[i:start])
        end = text.find("</think>", start + len("<think>"))
        if end == -1:
            # Unterminated think; treat remainder as thought
            thought_parts.append(text[start + len("<think>") :])
            break
        thought_parts.append(text[start + len("<think>") : end])
        i = end + len("</think>")
    visible = "".join(visible_parts)
    thought = "".join(thought_parts) if thought_parts else None
    return visible, thought


class OptimizerPapyrusClient:
    """
    Lightweight wrapper to call Papyrus for prompt optimization and persist reasoning output.

    The Papyrus client must already handle streaming and reasoning extraction; this wrapper
    saves any reasoning ("think" content) to the configured directory for later inspection.
    """

    def __init__(
        self,
        *,
        endpoint: str,
        model: str,
        token_provider: Callable[[], str] | None = None,
        bearer_token: str | None = None,
        thought_dir: str | os.PathLike[str],
        papyrus_quota_id: str = "",
        papyrus_timeout_ms: int = 100_000,
        request_timeout_s: float | None = None,
    ) -> None:
        self._thought_dir = Path(thought_dir)
        self._thought_dir.mkdir(parents=True, exist_ok=True)
 
        papyrus_client = build_papyrus_model_client(
            endpoint=endpoint,
            model=model,                  
            token_provider=token_provider,
            quota_id=papyrus_quota_id,
            papyrus_timeout_ms=papyrus_timeout_ms,
            request_timeout_s=request_timeout_s)
        self._client = papyrus_client

    async def close(self) -> None:
        await self._client.close()

    def _save_thought(self, experiment_id: str, thought: str) -> Path:
        path = self._thought_dir / f"{experiment_id}.txt"
        path.write_text(thought, encoding="utf-8")
        return path

    def extract_think_and_save(self, content: str, experiment_id: str) -> str:
        """
        Extract <think>...</think> sections from `content`, save them, and return the thought text.
        If no think sections are present, returns an empty string.
        """
        _, thought = _split_think_blocks(content)
        if thought:
            self._save_thought(experiment_id, thought)
            return thought
        return ""

    async def generate_prompt(self, user_query: str, experiment_id: str) -> str:
        """
        Call Papyrus to generate an optimized prompt for the given user query and experiment id.

        Returns the model's visible content (without <think> tags). Any reasoning is saved to `thought_dir`.
        """
        system = SystemMessage(content="You are a prompt optimizer. Improve prompts clearly and concisely.")
        user = UserMessage(content=user_query, source="optimizer")
        result: CreateResult = await self._client.create([system, user])

        # Prefer reasoning surfaced via CreateResult.thought; fall back to inline <think> blocks in content.
        thought_text = result.thought or ""
        content_text: str
        if isinstance(result.content, str):
            visible, inline_thought = _split_think_blocks(result.content)
            content_text = visible
            if not thought_text and inline_thought:
                thought_text = inline_thought
        else:
            content_text = str(result.content)

        if thought_text:
            self._save_thought(experiment_id, thought_text)

        return content_text
