import logging
import os
import re
import time
from io import TextIOWrapper
from json import JSONDecodeError
from typing import Dict, List, Optional, Tuple
from urllib.parse import urlparse

import openai
from hep_data_llm.utils import diskcache_decorator
from dotenv import dotenv_values, find_dotenv
from pydantic import BaseModel
from hep_data_llm.query_config import load_yaml_file
from tenacity import (
    before_sleep_log,
    retry,
    retry_if_exception_type,
    stop_after_attempt,
    wait_fixed,
)

from hep_data_llm.usage_info import UsageInfo, get_usage_info


class ModelInfo(BaseModel):
    model_name: str
    input_cost_per_million: float
    output_cost_per_million: float
    endpoint: Optional[str] = None  # e.g., OpenAI API endpoint or local server URL


def load_models(models_path: str = "models.yaml") -> Dict[str, ModelInfo]:
    """
    Load models and their costs from a YAML file, returning a dict of model_name to ModelInfo.
    """
    data = load_yaml_file(models_path)
    raw_models = data["models"]
    return {name: ModelInfo(**info) for name, info in raw_models.items()}


def process_model_request(
    requested_models: Optional[str],
    all_models: Dict[str, ModelInfo],
    default_model_name: str,
) -> List[str]:
    """
    Processes the requested model names and returns a validated list of model names.
    Built to be used with a command line option.

    Args:
        requested_models (Optional[str]): Comma-separated string of requested model names, or None.
        all_models (Dict[str, ModelInfo]): Dictionary of available models.
        default_model_name (str): The default model name to use if none are requested.

    Returns:
        List[str]: List of validated model names.

    Raises:
        ValueError: If any requested model name is not found in the available models.
    """
    if requested_models:
        model_names = [m.strip() for m in requested_models.split(",") if m.strip()]
        if "all" in model_names:
            model_names = list(all_models.keys())
    else:
        model_names = [default_model_name]

    # Validate model names
    invalid_model_names = [m for m in model_names if m not in all_models]
    if invalid_model_names:
        raise ValueError(
            f"Error: model(s) not found in models.yaml: {', '.join(invalid_model_names)}"
        )
    return model_names


@diskcache_decorator(".openai_response_cache")
def _get_openai_response(prompt: str, model_name: str, endpoint: Optional[str] = None):
    if endpoint:
        client = openai.OpenAI(base_url=endpoint)
    else:
        client = openai.OpenAI()

    @retry(
        stop=stop_after_attempt(3),
        wait=wait_fixed(0.5),
        retry=retry_if_exception_type(JSONDecodeError),
        before_sleep=before_sleep_log(logging.getLogger(__name__), logging.WARNING),
    )
    def do_request():
        start_time = time.time()
        response = client.chat.completions.create(
            model=model_name, messages=[{"role": "user", "content": prompt}]
        )
        elapsed = time.time() - start_time
        assert response.choices[0].message.content is not None, "No content in response"
        return {"response": response, "elapsed": elapsed}

    return do_request()


def run_llm(
    prompt: str, model_info: ModelInfo, out: TextIOWrapper, ignore_cache=False
) -> Tuple[UsageInfo, str]:
    # Set API key based on endpoint hostname, using <node-name>_API_KEY
    endpoint_host = None
    if model_info.endpoint:
        endpoint_host = urlparse(model_info.endpoint).hostname
    if not endpoint_host:
        endpoint_host = "api.openai.com"
    env_var = f"{endpoint_host.replace('.', '_')}_API_KEY"
    env_path = find_dotenv()
    env_vars = dotenv_values(env_path)
    api_key = env_vars.get(env_var)
    if api_key:
        os.environ["OPENAI_API_KEY"] = api_key
        openai.api_key = api_key
    else:
        logging.warning(f"API key not found for {env_var}")
        if "OPENAI_API_KEY" in env_vars:
            del os.environ["OPENAI_API_KEY"]

    # Do the query
    llm_result = _get_openai_response(
        prompt,
        model_info.model_name,
        model_info.endpoint,
        ignore_cache=ignore_cache,  # type: ignore
    )
    response = llm_result["response"]
    elapsed = llm_result["elapsed"]
    message = None
    if response and response.choices and response.choices[0].message:
        message = response.choices[0].message.content

    out.write("\n")
    if message:
        cleaned_message = strip_being_end(message)
        out.write(cleaned_message + "\n")
    else:
        out.write("No response content returned.")

    usage_info = get_usage_info(response, model_info, elapsed)

    return usage_info, str(message)


def strip_being_end(message: str) -> str:
    "Strip the being and end response strings off"
    message = message.strip()
    if message.startswith(">>start-reply<<"):
        message = message[15:]

    if message.endswith(">>end-reply<<"):
        message = message[: len(message) - 13]

    return message.strip()


def ensure_closing_triple_backtick(message: str) -> str:
    """
    Ensure that if a message contains an opening triple backtick, it also has a closing one.
    If the number of triple backticks is odd, append a closing triple backtick.
    """
    if "```" in message:
        backtick_count = message.count("```")
        if backtick_count % 2 != 0:
            message = message + "\n```"
    return message


def extract_code_from_response(message: str) -> Optional[str]:
    """
    Extract Python code from an OpenAI response object.
    Looks for code blocks in the message content and returns the first Python block
    found.
    """
    if not message:
        return None
    message = ensure_closing_triple_backtick(message)

    # Find all Python code blocks
    code_blocks = re.findall(r"```python(.*?)```", message, re.DOTALL | re.IGNORECASE)
    if code_blocks:
        if len(code_blocks) != 1:
            raise ValueError("Expected exactly one code block")
        return code_blocks[-1].strip()
    # Fallback: any code block
    code_blocks = re.findall(r"```(.*?)```", message, re.DOTALL)
    if code_blocks:
        return code_blocks[0].strip()
    return None


def extract_by_phase(text: str) -> Dict[str, str]:
    """
    Extracts sections from the input string that start with '## Phase XX',
    returning a dictionary mapping 'XX' to the section text (up to the next '##').
    """
    phase_pattern = re.compile(r"^## Phase ([^\n]+)", re.MULTILINE)
    matches = list(phase_pattern.finditer(text))
    result = {}
    for i, match in enumerate(matches):
        phase_name = match.group(1).strip()
        start = match.end()
        if i + 1 < len(matches):
            end = matches[i + 1].start()
        else:
            end = len(text)
        section = text[start:end].strip()
        result[phase_name] = section
    return result
