import asyncio
import json
import os
from typing import Any, Literal, cast

from json_schema_to_pydantic import create_model as create_pydantic_model_from_json_schema
from langchain_core.messages import HumanMessage
from loguru import logger
from pydantic import BaseModel, Field
from universal_mcp.applications.application import BaseApplication

from universal_mcp.agents.llm import load_chat_model
from universal_mcp.agents.utils import get_message_text

MAX_RETRIES = 3


def _get_context_as_string(source: Any | list[Any] | dict[str, Any]) -> str:
    """Converts context to a string representation.

    Args:
        source: The source data to be converted. Can be a single value, a list of values, or a dictionary.

    Returns:
        A string representation of the source data, formatted with XML-like tags for dictionaries.
    """


def _format_image_content(images: list[Any]) -> list[dict[str, Any]]:
    """Format a list of images (URLs, dicts, etc.) into LangChain image blocks."""
    formatted_images = []
    for img in images:
        if isinstance(img, dict) and img.get("type") == "image_url":
            formatted_images.append(img)
        elif isinstance(img, dict) and "url" in img:
            formatted_images.append({"type": "image_url", "image_url": {"url": img["url"]}})
        elif isinstance(img, str):  # direct URL string
            formatted_images.append({"type": "image_url", "image_url": {"url": img}})
        # Try to include raw dicts if they fit the schema, otherwise skip/warn?
        # For now, just append if it looks dict-like, assume user knows schema
        elif isinstance(img, dict):
            formatted_images.append(img)
    return formatted_images


def _get_context_as_string(source: Any | list[Any] | dict[str, Any]) -> str:
    """Converts a context object to a string representation."""
    if not isinstance(source, dict):
        if isinstance(source, list):
            source = {f"doc_{i + 1}": doc for i, doc in enumerate(source)}
        else:
            source = {"content": source}

    return "\n".join(f"<{k}>\n{str(v)}\n</{k}>" for k, v in source.items())


def _get_default_llm_name() -> str:
    """Return LLM model name from env with fallback."""
    return os.getenv("LLM_TOOLS_LLM", "anthropic:claude-haiku-4-5")


def _as_dict(response: Any) -> dict[str, Any]:
    """Normalize structured output responses across providers."""
    if isinstance(response, dict):
        return response
    if hasattr(response, "model_dump"):
        return cast(dict[str, Any], response.model_dump())
    return cast(dict[str, Any], response)


class LlmApp(BaseApplication):
    """
    An application for leveraging Large Language Models (LLMs) for advanced text processing tasks.
    """

    def __init__(self, **kwargs):
        """Initialize the LLMApp."""
        super().__init__(name="llm")

    async def generate_text(
        self,
        task: str = "",
        context: str | list[str] | dict[str, str] = "",
        tone: str = "normal",
        output_format: Literal["markdown", "html", "plain"] = "markdown",
        length: Literal["very-short", "concise", "normal", "long"] = "concise",
        images: list[Any] | None = None,
        temperature: float = 0.3,
    ) -> str:
        """
        Given a high-level writing task and context, returns a well-written text
        that achieves the task, given the context.

        Example Call:
            generate_text(task="Summarize this website with the goal of making it easy to understand.", context=web_content)
            generate_text(task="Create an HTML dashboard", output_format="html", context=web_content)
            generate_text(task="Prepare a markdown report based on this content", output_format="markdown", context=web_content)
            generate_text(task="Make a markdown table summarizing the key differences between doc_1 and doc_2.", context={"doc_1": str(doc_1), "doc_2": str(doc_2)})
            generate_text(task="Summarize all the provided documents.", context=[doc_1, doc_2, doc_3])

        Important:
        - Include specifics of the goal in the context verbatim.
        - Be precise and direct in the task, and include as much context as possible.
        - Include relevant high-level goals or intent in the task.
        - You can provide multiple documents as input, and reference them in the task.
        - You MUST provide the contents of any source documents to `generate_text`.
        - NEVER use `generate_text` to produce JSON for a Pydantic model.

        Args:
            task: The main writing task or directive. Optional if images are provided.
            context: A single string, list of strings, or dict mapping labels to content.
            tone: The desired tone of the output (e.g., "normal", "flirty", "formal", "casual", "crisp", "poetic", "technical", "internet-chat", "smartass", etc.).
            output_format: The desired output format ('markdown', 'html', 'plain').
            length: Desired length of the output ('very-short', 'concise', 'normal', 'long').
            images: Optional list of images. Each item can be:
                - A string URL (e.g. "https://example.com/image.jpg")
                - A dict with a "url" key (e.g. {"url": "..."})
                - A LangChain image block (e.g. {"type": "image_url", "image_url": {"url": "..."}})
            temperature: The temperature for the LLM (0-1). Higher values for more creativity, lower for precision. Defaults to 0.3.

        Returns:
            The generated text as a string.
        """
        context_str = _get_context_as_string(context)

        prompt = f"{task.strip()}\n\n"
        if output_format == "markdown":
            prompt += "Please write in Markdown format.\n\n"
        elif output_format == "html":
            prompt += "Please write in HTML format.\n\n"
        else:
            prompt += "Please write in plain text format. Do not use markdown or HTML.\n\n"

        if tone not in ["normal", "default", ""]:
            prompt = f"{prompt} (Tone instructions: {tone})"

        if length not in ["normal", "default", ""]:
            prompt = f"{prompt} (Length instructions: {length})"

        full_prompt_text = f"{prompt}\n\nContext:\n{context_str}\n\n"

        message_content = [{"type": "text", "text": full_prompt_text}]
        if images:
            image_blocks = _format_image_content(images)
            message_content.extend(image_blocks)

        final_prompt = HumanMessage(content=message_content) if images else full_prompt_text

        model = load_chat_model(
            _get_default_llm_name(),
            temperature=temperature,
            disable_streaming=True,
            thinking=False,
            streaming=False,
            tags=("quiet",),
        )
        response = (
            await model.with_retry(stop_after_attempt=MAX_RETRIES)
            .with_config({"tags": ["nostream"]})
            .ainvoke([final_prompt] if isinstance(final_prompt, HumanMessage) else final_prompt)
        )
        return get_message_text(response)

    async def classify_data(
        self,
        class_descriptions: dict[str, str],
        classification_task_and_requirements: str = "",
        context: Any | list[Any] | dict[str, Any] = "",
        images: list[Any] | None = None,
    ) -> dict[str, Any]:
        """
        Classifies and compares data based on given requirements.

        Use `classify_data` for tasks which need to classify data into one of many categories.
        If making multiple binary classifications, call `classify_data` for each.

        Guidance:
        - Prefer to use classify_data operations to compare strings, rather than string ops.
        - Prefer to include an "Unsure" category for classification tasks.
        - The `class_descriptions` dict argument MUST be a map from possible class names to a precise description.
        - Use precise and specific class names and concise descriptions.
        - Pass ALL relevant context, preferably as a dict mapping labels to content.
        - Returned dict maps each possible class name to a probability.

        Example Usage:
            classification_task_and_requirements = "Does the document contain an address?"
            class_descriptions = {
                "Is_Address": "Valid addresses usually have street names, city, and zip codes.",
                "Not_Address": "Not valid addresses."
            }
            classification = classify_data(
                classification_task_and_requirements=classification_task_and_requirements,
                context={"address": extracted_address},
                class_descriptions=class_descriptions
            )
            if classification["probabilities"]["Is_Address"] > 0.5:
                ...

        Args:
            class_descriptions: Mapping from class names to descriptions.
            classification_task_and_requirements: The classification question and rules. Optional if images provided.
            context: The data to classify (string, list, or dict).
            images: Optional list of images. Each item can be:
                - A string URL (e.g. "https://example.com/image.jpg")
                - A dict with a "url" key (e.g. {"url": "..."})
                - A LangChain image block (e.g. {"type": "image_url", "image_url": {"url": "..."}})

        Tags:
            important

        Returns:
            dict: {
                "reason": str,
                "top_class": str,
            }
        """
        context_str = _get_context_as_string(context)

        prompt = (
            f"{classification_task_and_requirements}\n\n"
            f"This is a classification task.\nPossible classes and descriptions:\n"
            f"{json.dumps(class_descriptions, indent=2)}\n\n"
            f"Context:\n{context_str}\n\n"
            "Return ONLY a valid JSON object that conforms to the provided schema, with no extra text."
        )

        message_content = [{"type": "text", "text": prompt}]
        if images:
            image_blocks = _format_image_content(images)
            message_content.extend(image_blocks)
        final_prompt = HumanMessage(content=message_content) if images else prompt

        class ClassificationResult(BaseModel):
            reason: str = Field(..., description="The reasoning behind the classification.")
            top_class: str = Field(..., description="The class with the highest probability.")
            # probabilities: dict[str, float] = Field(..., description="The probabilities for each class.")

        model = load_chat_model(
            _get_default_llm_name(),
            temperature=0.3,
            disable_streaming=True,
            thinking=False,
            tags=("quiet",),
            streaming=False,
        )
        response = await (
            model.with_structured_output(ClassificationResult)
            .with_retry(stop_after_attempt=MAX_RETRIES)
            .with_config({"tags": ["nostream"]})
            .ainvoke([final_prompt] if isinstance(final_prompt, HumanMessage) else final_prompt)
        )
        return _as_dict(response)

    async def extract_data(
        self,
        output_schema: dict[str, Any],
        extraction_task: str = "",
        source: Any | list[Any] | dict[str, Any] = "",
        images: list[Any] | None = None,
    ) -> dict[str, Any]:
        """
        Extracts structured data from unstructured data (documents, webpages, images, large bodies of text),
        returning a dictionary matching the given output_schema.

        You MUST anticipate Exception raised for unextractable data; skip this item if applicable.

        Strongly prefer to:
        - Be comprehensive, specific, and precise on the data you want to extract.
        - Use optional fields everywhere.
        - Extract multiple items from each source unless otherwise specified.
        - The more specific your extraction task and output_schema are, the better the results.

        Args:
            output_schema: must be a valid JSON schema with top-level 'title' and 'description' keys. Title must be a safe identifier (slug).
            extraction_task: The directive describing what to extract. Optional if images provided.
            source: The unstructured data to extract from.
            images: Optional list of images. Each item can be:
                - A string URL (e.g. "https://example.com/image.jpg")
                - A dict with a "url" key (e.g. {"url": "..."})
                - A LangChain image block (e.g. {"type": "image_url", "image_url": {"url": "..."}})

        Returns:
            A dictionary containing the extracted data, matching the provided schema.

        Example:
            news_articles_schema = {
                "title": "News_Article_List",
                "description": "A list of news articles with headlines and URLs",
                "type": "object",
                "properties": {
                    "articles": {
                        "type": "array",
                        "items": {
                            "type": "object",
                            "properties": {
                                "headline": {
                                    "type": "string"
                                },
                                "url": {
                                    "type": "string"
                                }
                            },
                            "required": ["headline", "url"]
                        }
                    }
                },
                "required": ["articles"]
            }

            news_articles = extract_data(extraction_task="Extract headlines and their corresponding URLs.", source=content, output_schema=news_articles_schema)
        """
        context_str = _get_context_as_string(source)

        prompt = (
            f"{extraction_task}\n\n"
            f"Context:\n{context_str}\n\n"
            "Return ONLY a valid JSON object that conforms to the provided schema, with no extra text."
        )

        message_content = [{"type": "text", "text": prompt}]
        if images:
            image_blocks = _format_image_content(images)
            message_content.extend(image_blocks)
        final_prompt = HumanMessage(content=message_content) if images else prompt

        model = load_chat_model(
            _get_default_llm_name(),
            temperature=0.3,
            disable_streaming=True,
            thinking=False,
            tags=("quiet",),
            streaming=False,
        )

        PModel = create_pydantic_model_from_json_schema(output_schema)
        response = (
            await model.with_structured_output(PModel)
            .with_retry(stop_after_attempt=MAX_RETRIES)
            .with_config({"tags": ["nostream"]})
            .ainvoke([final_prompt] if isinstance(final_prompt, HumanMessage) else final_prompt)
        )
        return _as_dict(response)

    async def call_llm(
        self,
        output_schema: dict[str, Any],
        task_instructions: str = "",
        context: Any | list[Any] | dict[str, Any] = "",
        images: list[Any] | None = None,
        temperature: float = 0.3,
    ) -> dict[str, Any]:
        """
        Call a Large Language Model (LLM) with an instruction and contextual information,
        returning a dictionary matching the given output_schema.
        Can be used for tasks like creative writing, llm reasoning based content generation, etc.

        You MUST anticipate Exceptions in reasoning based tasks which will lead to some empty fields
        in the returned output; skip this item if applicable.

        General Guidelines:
        - Be comprehensive, specific, and precise on the task instructions.
        - Include as much context as possible.
        - You can provide multiple items in context, and reference them in the task.
        - Include relevant high-level goals or intent in the task.
        - In the output_schema, use required field wherever necessary.
        - The more specific your task instructions and output_schema are, the better the results.

        Guidelines for content generation tasks:
        - Feel free to add instructions for tone, length, and format (markdown, html, plain-text, xml)
        - Some examples of tone are: "normal", "flirty", "formal", "casual", "crisp", "poetic", "technical", "internet-chat", "smartass", etc.
        - Prefer length to be concise by default. Other examples are: "very-short", "concise", "normal", "long", "2-3 lines", etc.
        - In format prefer plain-text but you can also use markdown and html wherever useful.

        Args:
            output_schema: must be a valid JSON schema with top-level 'title' and 'description' keys. Title must be a safe identifier (slug).
            task_instructions: The main directive for the LLM (e.g., "Summarize the article" or "Extract key entities"). Optional if images provided.
            context:
                A dictionary containing named text elements that provide additional
                information for the LLM. Keys are labels (e.g., 'article', 'transcript'),
                values are strings of content.
            images: Optional list of images. Each item can be:
                - A string URL (e.g. "https://example.com/image.jpg")
                - A dict with a "url" key (e.g. {"url": "..."})
                - A LangChain image block (e.g. {"type": "image_url", "image_url": {"url": "..."}})
            temperature: The temperature for the LLM (0-1). Higher values for more creativity, lower for precision. Defaults to 0.3.

        Returns:
            dict: Parsed JSON object matching the desired output_schema.

        """
        context_str = _get_context_as_string(context)

        prompt = f"{task_instructions}\n\nContext:\n{context_str}\n\nReturn ONLY a valid JSON object, no extra text."

        message_content = [{"type": "text", "text": prompt}]
        if images:
            image_blocks = _format_image_content(images)
            message_content.extend(image_blocks)
        final_prompt = HumanMessage(content=message_content) if images else prompt

        model = load_chat_model(
            _get_default_llm_name(),
            temperature=temperature,
            disable_streaming=True,
            thinking=False,
            tags=("quiet",),
            streaming=False,
        )

        PModel = create_pydantic_model_from_json_schema(output_schema)
        model_with_structure = model.with_structured_output(PModel)
        response = (
            await model_with_structure.with_retry(stop_after_attempt=MAX_RETRIES)
            .with_config({"tags": ["nostream"]})
            .ainvoke([final_prompt] if isinstance(final_prompt, HumanMessage) else final_prompt)
        )
        return _as_dict(response)

    def list_tools(self):
        return [
            self.generate_text,
            self.classify_data,
            self.extract_data,
            self.call_llm,
        ]


async def test():
    """Exercise each tool once."""
    app = LlmApp()

    generated = await app.generate_text(task="Say hello to the world", output_format="plain")
    logger.info({"generate_text": generated})

    classification = await app.classify_data(
        classification_task_and_requirements="Does the document contain an address?",
        class_descriptions={
            "Is_Address": "Valid addresses usually have street names, city, and zip codes.",
            "Not_Address": "Not valid addresses.",
        },
        context={"address": "123 Main St, Anytown, USA"},
    )
    logger.info({"classify_data": classification})

    extract_schema = {
        "title": "AddressExtraction",
        "description": "Extract address line if present",
        "type": "object",
        "properties": {"address": {"type": "string"}},
        "required": ["address"],
    }
    extracted = await app.extract_data(
        output_schema=extract_schema, extraction_task="Extract the address", source={"doc": "123 Main St"}
    )
    logger.info({"extract_data": extracted})

    call_schema = {
        "title": "Greeting",
        "description": "Simple greeting schema",
        "type": "object",
        "properties": {"greeting": {"type": "string"}},
        "required": ["greeting"],
    }
    llm_output = await app.call_llm(
        output_schema=call_schema,
        task_instructions="Return a short greeting.",
        context={"context": "No extra context"},
    )
    logger.info({"call_llm": llm_output})

    # Test Image Input
    image_url = "https://upload.wikimedia.org/wikipedia/commons/thumb/d/dd/Gfp-wisconsin-madison-the-nature-boardwalk.jpg/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
    image_task = "Describe this image in detail."
    image_response = await app.generate_text(task=image_task, images=[image_url])
    logger.info({"generate_text_image": image_response})


if __name__ == "__main__":
    asyncio.run(test())
