import base64
import ipaddress
import json
import logging
import mimetypes
import socket
from collections.abc import Callable
from datetime import date, datetime
from decimal import Decimal
from enum import Enum
from pathlib import Path
from typing import Any, Literal, Optional, Union
from urllib.parse import urlparse

import httpx
import jsonref
from agents import FunctionTool, ToolOutputFileContent, ToolOutputImage
from agents.run_context import RunContextWrapper
from agents.strict_schema import ensure_strict_json_schema
from datamodel_code_generator import DataModelType, PythonVersion
from datamodel_code_generator.model import get_data_model_types
from datamodel_code_generator.parser.jsonschema import JsonSchemaParser

logger = logging.getLogger(__name__)

PDF_MIME_TYPE = "application/pdf"
URL_FETCH_TIMEOUT_SECONDS = 20.0
MAX_INLINE_PDF_BYTES = 10 * 1024 * 1024
MAX_FETCH_REDIRECTS = 5
REDIRECT_STATUS_CODES = {301, 302, 303, 307, 308}


def _build_data_url_from_bytes(data: bytes, mime_type: str) -> str:
    encoded_file = base64.b64encode(data).decode("utf-8")
    return f"data:{mime_type};base64,{encoded_file}"


def _build_data_url(file_path: Path, mime_type: str) -> str:
    return _build_data_url_from_bytes(file_path.read_bytes(), mime_type)


def _resolve_mime_type(file_path: Path) -> str:
    mime_type, _ = mimetypes.guess_type(file_path.name)
    if not mime_type:
        raise ValueError(f"Unable to determine MIME type for file: {file_path}")
    return mime_type


def _extract_filename_from_url(url: str) -> str | None:
    try:
        parsed = urlparse(url)
    except ValueError:
        return None
    filename = Path(parsed.path).name
    return filename or None


def _is_global_ip_address(value: str) -> bool:
    try:
        ip = ipaddress.ip_address(value)
    except ValueError:
        return False
    return ip.is_global


def _is_remote_host_safe_for_fetch(url: str) -> bool:
    try:
        parsed = urlparse(url)
    except ValueError:
        return False
    if parsed.scheme not in {"http", "https"}:
        return False

    try:
        parsed_port = parsed.port
    except ValueError:
        return False

    hostname = parsed.hostname
    if not hostname:
        return False
    if hostname.lower() in {"localhost", "localhost.localdomain"}:
        return False

    # Direct IPs are validated without DNS lookups.
    if _is_global_ip_address(hostname):
        return True

    if _is_global_ip_address(hostname.strip("[]")):
        return True

    port = parsed_port or (443 if parsed.scheme == "https" else 80)
    try:
        addr_infos = socket.getaddrinfo(hostname, port, type=socket.SOCK_STREAM)
    except socket.gaierror:
        return False

    resolved_ips: set[str] = set()
    for entry in addr_infos:
        socket_address = entry[4]
        if not isinstance(socket_address, tuple) or not socket_address:
            continue
        host = socket_address[0]
        if isinstance(host, str):
            resolved_ips.add(host)
    if not resolved_ips:
        return False

    return all(_is_global_ip_address(ip.strip("[]")) for ip in resolved_ips)


def _is_pdf_content_type(content_type: str | None) -> bool:
    if not content_type:
        return False
    normalized_content_type = content_type.split(";", 1)[0].strip().lower()
    return normalized_content_type == PDF_MIME_TYPE


def _resolve_redirect_target(response: httpx.Response) -> str:
    location = response.headers.get("location")
    if not location:
        raise ValueError("Redirect response missing Location header")
    request_url = getattr(response, "request", None)
    if request_url is None:
        raise ValueError("Redirect response missing request URL")
    return str(request_url.url.join(location))


def _fetch_remote_headers(url: str) -> httpx.Headers:
    current_url = url
    for _ in range(MAX_FETCH_REDIRECTS + 1):
        response = httpx.head(current_url, follow_redirects=False, timeout=URL_FETCH_TIMEOUT_SECONDS)
        if response.status_code in REDIRECT_STATUS_CODES:
            next_url = _resolve_redirect_target(response)
            if not _is_remote_host_safe_for_fetch(next_url):
                raise ValueError(f"Unsafe redirect target: {next_url}")
            current_url = next_url
            continue
        response.raise_for_status()
        return response.headers

    raise ValueError(f"Exceeded max redirects ({MAX_FETCH_REDIRECTS}) while fetching headers: {url}")


def _download_with_size_limit(url: str, *, max_bytes: int) -> bytes:
    current_url = url
    for _ in range(MAX_FETCH_REDIRECTS + 1):
        buffer = bytearray()
        with httpx.stream("GET", current_url, follow_redirects=False, timeout=URL_FETCH_TIMEOUT_SECONDS) as response:
            if response.status_code in REDIRECT_STATUS_CODES:
                next_url = _resolve_redirect_target(response)
                if not _is_remote_host_safe_for_fetch(next_url):
                    raise ValueError(f"Unsafe redirect target: {next_url}")
                current_url = next_url
                continue

            response.raise_for_status()
            for chunk in response.iter_bytes():
                if not chunk:
                    continue
                buffer.extend(chunk)
                if len(buffer) > max_bytes:
                    raise ValueError(f"Remote file exceeds max inline size ({max_bytes} bytes)")
            return bytes(buffer)

    raise ValueError(f"Exceeded max redirects ({MAX_FETCH_REDIRECTS}) while downloading: {url}")


def _inline_pdf_url_as_file_data(url: str, *, filename: str) -> ToolOutputFileContent:
    content = _download_with_size_limit(url, max_bytes=MAX_INLINE_PDF_BYTES)
    return ToolOutputFileContent(
        file_data=_build_data_url_from_bytes(content, PDF_MIME_TYPE),
        filename=filename,
    )


def tool_output_image_from_path(
    path: str | Path,
    *,
    detail: Literal["auto", "high", "low"] = "auto",
) -> ToolOutputImage:
    """
    Build a ``ToolOutputImage`` from a local image file by returning a data URL.

    Args:
        path: Path to the image file on disk.
        detail: Optional detail hint to forward to the vision model.

    Raises:
        ValueError: If the file type cannot be resolved from the path.
    """

    file_path = Path(path)
    mime_type = _resolve_mime_type(file_path)
    return ToolOutputImage(image_url=_build_data_url(file_path, mime_type), detail=detail)


def tool_output_image_from_file_id(
    file_id: str,
    *,
    detail: Literal["auto", "high", "low"] = "auto",
) -> ToolOutputImage:
    """
    Build a ``ToolOutputImage`` from an OpenAI file ID.

    Args:
        file_id: openai file id of the image file.
        detail: Optional detail hint to forward to the vision model.
    """

    return ToolOutputImage(file_id=file_id, detail=detail)


def tool_output_file_from_path(path: str | Path, *, filename: str | None = None) -> ToolOutputFileContent:
    """
    Build a ``ToolOutputFileContent`` from a local file by embedding base64 data.

    Args:
        path: Path to the file on disk.
        filename: Optional filename hint for the client.

    Raises:
        ValueError: If the file is not a PDF.
    """

    file_path = Path(path)
    if filename and not filename.lower().endswith(".pdf"):
        raise ValueError(f"Filename must end with .pdf, got: {filename}")
    mime_type = _resolve_mime_type(file_path)
    if mime_type != PDF_MIME_TYPE:
        raise ValueError("Only PDF files are supported.")
    return ToolOutputFileContent(
        file_data=_build_data_url(file_path, PDF_MIME_TYPE), filename=filename or file_path.name
    )


def tool_output_file_from_url(url: str) -> ToolOutputFileContent:
    """
    Build a ``ToolOutputFileContent`` that references an externally hosted file.

    Args:
        url: Publicly reachable URL for the file.

    Notes:
        For ``.pdf`` URLs, this helper first checks the remote ``Content-Type``.
        If the host reports a non-PDF type (for example ``application/octet-stream``),
        it falls back to downloading and inlining the bytes as ``file_data`` with
        ``application/pdf`` MIME to keep Responses API ingestion reliable.
    """
    filename = _extract_filename_from_url(url)
    if not filename or not filename.lower().endswith(".pdf"):
        return ToolOutputFileContent(file_url=url)

    if not _is_remote_host_safe_for_fetch(url):
        logger.warning("Skipping local fetch for potentially unsafe remote PDF URL %s.", url)
        return ToolOutputFileContent(file_url=url)

    try:
        response_headers = _fetch_remote_headers(url)
    except Exception:
        logger.debug("Failed to inspect remote file URL %s; preserving file_url behavior.", url, exc_info=True)
        return ToolOutputFileContent(file_url=url)

    if _is_pdf_content_type(response_headers.get("content-type")):
        return ToolOutputFileContent(file_url=url)

    try:
        return _inline_pdf_url_as_file_data(url, filename=filename)
    except Exception:
        logger.warning(
            "Failed to inline PDF URL %s after non-PDF content-type (%s); preserving file_url behavior.",
            url,
            response_headers.get("content-type"),
            exc_info=True,
        )
        return ToolOutputFileContent(file_url=url)


def tool_output_file_from_file_id(file_id: str) -> ToolOutputFileContent:
    """
    Build a ``ToolOutputFileContent`` that references an openai file id.

    Args:
        file_id: openai file id of the pdf file.
    """

    return ToolOutputFileContent(file_id=file_id)


def from_openapi_schema(
    schema: str | dict[str, Any],
    *,
    headers: dict[str, str] | None = None,
    params: dict[str, Any] | None = None,
    strict: bool = False,
    timeout: int = 90,
) -> list[FunctionTool]:
    """
    Converts an OpenAPI JSON or dictionary describing a single endpoint into one or more FunctionTool instances.

    Args:
        schema (str | dict): Full OpenAPI JSON string or dictionary.
        headers (dict[str, str] | None, optional): Extra HTTP headers to send with each call. Defaults to None.
        params (dict[str, Any] | None, optional): Extra query parameters to append to every call. Defaults to None.
        strict (bool, optional): If True, sets 'additionalProperties' to False in every generated schema.
            Defaults to False.
        timeout (int, optional): HTTP timeout in seconds. Defaults to 90.

    Returns:
        list[FunctionTool]: List of FunctionTool instances generated from the OpenAPI endpoint.
    """

    if isinstance(schema, dict):
        openapi = jsonref.JsonRef.replace_refs(schema)
    else:
        openapi = jsonref.loads(schema)

    headers = {k: v for k, v in (headers or {}).items() if v is not None}
    fixed_params = params or {}

    tools: list[FunctionTool] = []

    base_url = openapi["servers"][0]["url"]

    for path, verbs in openapi["paths"].items():
        for verb, verb_spec_ref in verbs.items():
            verb_spec = jsonref.replace_refs(verb_spec_ref)

            # Build OpenAI-compatible JSON schema

            function_name = verb_spec.get("operationId")
            description = verb_spec.get("description") or verb_spec.get("summary", "")

            req_body_schema = (
                verb_spec.get("requestBody", {}).get("content", {}).get("application/json", {}).get("schema")
            )

            parameters_obj_schema = build_parameter_object_schema(
                verb_spec.get("parameters", []),
                strict,
            )

            tool_schema = build_tool_schema(parameters_obj_schema, req_body_schema, strict=strict)

            # Callback factory  (captures current verb & path)

            async def _invoke(
                ctx: RunContextWrapper[Any],
                input_json: str,
                *,
                verb_: str = verb,
                path_: str = path,
                base_url_: str = base_url,
            ):
                """Actual HTTP call executed by the agent."""
                payload = json.loads(input_json)

                # split out parts for old-style structure
                raw_parameters: dict[str, Any] = payload.get("parameters", {})
                body_payload = payload.get("requestBody")

                url, remaining_params = resolve_url(base_url_, path_, raw_parameters)

                query_params = {k: v for k, v in remaining_params.items() if v is not None}
                if fixed_params:
                    query_params = {**query_params, **fixed_params}

                json_body = body_payload if verb_.lower() in {"post", "put", "patch", "delete"} else None

                logger.info(f"Calling URL: {url}\nQuery Params: {query_params}\nJSON Body: {json_body}")

                async with httpx.AsyncClient(timeout=timeout) as client:
                    resp = await client.request(
                        verb_.upper(),
                        url,
                        params=query_params,
                        json=json_body,
                        headers=headers,
                    )
                    try:
                        logger.info(f"Response from {url}: {resp.json()}")
                        return resp.json()
                    except Exception:
                        return resp.text

            tool = FunctionTool(
                name=function_name,
                description=description,
                params_json_schema=tool_schema,
                on_invoke_tool=_invoke,
                strict_json_schema=strict,
            )
            tools.append(tool)

    return tools


def validate_openapi_spec(spec: str):
    spec_dict = json.loads(spec)

    # Validate that 'paths' is present in the spec
    if "paths" not in spec_dict:
        raise ValueError("The spec must contain 'paths'.")

    paths = spec_dict["paths"]
    if not isinstance(paths, dict):
        raise ValueError("The 'paths' field must be a dictionary.")

    for path, path_item in paths.items():
        # Check that each path item is a dictionary
        if not isinstance(path_item, dict):
            raise ValueError(f"Path item for '{path}' must be a dictionary.")

        for operation in path_item.values():
            # Basic validation for each operation
            if "operationId" not in operation:
                raise ValueError("Each operation must contain an 'operationId'.")
            if "description" not in operation:
                raise ValueError("Each operation must contain a 'description'.")

    return spec_dict


def generate_model_from_schema(schema: dict, class_name: str, strict: bool) -> type:
    data_model_types = get_data_model_types(
        DataModelType.PydanticV2BaseModel,
        target_python_version=PythonVersion.PY_310,
    )
    parser = JsonSchemaParser(
        json.dumps(schema),
        data_model_type=data_model_types.data_model,
        data_model_root_type=data_model_types.root_model,
        data_model_field_type=data_model_types.field_model,
        data_type_manager_type=data_model_types.data_type_manager,
        dump_resolve_reference_action=data_model_types.dump_resolve_reference_action,
        use_schema_description=True,
        validation=False,
        class_name=class_name,
        strip_default_none=strict,
    )
    result = parser.parse()
    imports_str = "from typing import List, Dict, Any, Optional, Union, Set, Tuple, Literal\nfrom enum import Enum\n"
    if isinstance(result, str):
        result = imports_str + result
    else:
        result = imports_str + str(result)
    result = result.replace("from __future__ import annotations\n", "")
    result += f"\n\n{class_name}.model_rebuild(force=True)"
    exec_globals = {
        "List": list,
        "Dict": dict,
        "Type": type,
        "Union": Union,
        "Optional": Optional,
        "datetime": datetime,
        "date": date,
        "Set": set,
        "Tuple": tuple,
        "Any": Any,
        "Callable": Callable,
        "Decimal": Decimal,
        "Literal": Literal,
        "Enum": Enum,
    }
    exec(result, exec_globals)
    model = exec_globals.get(class_name)
    if not model:
        raise ValueError(f"Could not extract model from schema {class_name}")
    if hasattr(model, "model_rebuild"):
        try:
            model.model_rebuild(force=True)
        except Exception as e:
            print(f"Warning: Could not rebuild model {class_name} after exec: {e}")
    return model  # type: ignore[return-value]


def collect_parameter_schemas(parameters: list[dict[str, Any]]) -> tuple[dict[str, Any], list[str]]:
    """Extract parameter schemas and required flags from an OpenAPI operation."""
    properties: dict[str, Any] = {}
    required: list[str] = []

    for parameter in parameters:
        if "schema" not in parameter and "type" in parameter:
            parameter["schema"] = {"type": parameter["type"]}

        schema = parameter.get("schema")
        if not schema:
            raise ValueError(f"Parameter '{parameter['name']}' must define a schema")

        property_schema = properties.setdefault(parameter["name"], schema.copy())

        for attribute in ("description", "example", "examples"):
            if attribute in parameter:
                property_schema[attribute] = parameter[attribute]

        if parameter.get("required"):
            required.append(parameter["name"])

    return properties, required


def build_parameter_object_schema(parameters: list[dict[str, Any]], strict: bool) -> dict[str, Any]:
    properties, required = collect_parameter_schemas(parameters)
    return {
        "type": "object",
        "properties": properties,
        "required": required,
        "additionalProperties": False if strict else True,
    }


def build_tool_schema(
    parameter_schema: dict[str, Any],
    request_body_schema: dict[str, Any] | None,
    *,
    strict: bool,
    include_strict_flag: bool = False,
) -> dict[str, Any]:
    schema: dict[str, Any] = {
        "type": "object",
        "properties": {"parameters": parameter_schema},
        "required": ["parameters"],
        "additionalProperties": False if strict else True,
    }

    if include_strict_flag and strict:
        schema["strict"] = True

    if request_body_schema:
        body_schema = request_body_schema.copy()
        if strict:
            body_schema.setdefault("additionalProperties", False)
        schema["properties"]["requestBody"] = body_schema
        schema["required"].append("requestBody")

    return ensure_strict_json_schema(schema) if strict else schema


def resolve_url(base_url: str, path: str, parameters: dict[str, Any]) -> tuple[str, dict[str, Any]]:
    """Replace templated path parameters and return remaining query parameters."""
    url = f"{base_url}{path}"
    remaining_params: dict[str, Any] = {}

    for key, value in parameters.items():
        token = f"{{{key}}}"
        if token in url:
            url = url.replace(token, str(value))
            continue
        remaining_params[key] = value

    return url.rstrip("/"), remaining_params
