"""FastAPI scaffolding for Task Apps (local dev + deployment).

Prefer synth_ai.sdk.localapi.server moving forward. This module remains for
backward compatibility during the naming transition.
"""

from __future__ import annotations

import asyncio
import inspect
import os
import threading
from collections.abc import Awaitable, Callable, Iterable, Mapping, MutableMapping, Sequence
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import httpx
from fastapi import APIRouter, Depends, FastAPI, Query, Request
from fastapi.middleware.cors import CORSMiddleware
from starlette.middleware import Middleware

from synth_ai.core.telemetry import log_info

from .auth import normalize_environment_api_key, require_api_key_dependency
from .contracts import RolloutRequest, RolloutResponse, TaskInfo
from .datasets import TaskDatasetRegistry
from .errors import http_exception
from .json import to_jsonable
from .proxy import (
    inject_system_hint,
    prepare_for_groq,
    prepare_for_openai,
    synthesize_tool_call_if_missing,
)
from .rubrics import Rubric
from .vendors import get_groq_key_or_503, get_openai_key_or_503, normalize_vendor_keys

TasksetDescriptor = Callable[[], Mapping[str, Any] | Awaitable[Mapping[str, Any]]]
InstanceProvider = Callable[[Sequence[int]], Iterable[TaskInfo] | Awaitable[Iterable[TaskInfo]]]
RolloutExecutor = Callable[[RolloutRequest, Request], Any | Awaitable[Any]]


def _default_app_state() -> dict[str, Any]:
    return {}


@dataclass(slots=True)
class RubricBundle:
    """Optional rubrics advertised by the task app."""

    outcome: Rubric | None = None
    events: Rubric | None = None


@dataclass(slots=True)
class ProxyConfig:
    """Configuration for optional vendor proxy endpoints."""

    enable_openai: bool = False
    enable_groq: bool = False
    system_hint: str | None = None
    openai_url: str = "https://api.openai.com/v1/chat/completions"
    groq_url: str = "https://api.groq.com/openai/v1/chat/completions"


@dataclass(slots=True)
class TaskAppConfig:
    """Declarative configuration describing a Task App.
    
    A TaskAppConfig defines all the components needed to create a task app:
    - Task information (base_task_info)
    - Task set description (provide_taskset_description)
    - Task instance provider (provide_task_instances)
    - Rollout executor (rollout)
    - Optional rubrics, datasets, proxy config, etc.
    
    This config is used by `create_task_app()` to build a FastAPI application
    that implements the Synth AI task app contract.
    
    Example:
        >>> from synth_ai.sdk.task.server import TaskAppConfig, create_task_app
        >>>
        >>> def build_config() -> TaskAppConfig:
        ...     return TaskAppConfig(
        ...         app_id="my_task",
        ...         name="My Task App",
        ...         description="A simple task app",
        ...         provide_taskset_description=lambda: {"splits": ["train", "val"]},
        ...         provide_task_instances=lambda seeds: [...],
        ...         rollout=lambda req, r: {...},
        ...         # base_task_info is optional - auto-derived from app_id/name
        ...     )
        >>>
        >>> app = create_task_app(build_config())
        >>> # app is a FastAPI instance ready to run
    
    Attributes:
        app_id: Unique identifier for this task app
        name: Human-readable name
        description: Description of what this task app does
        base_task_info: Base TaskInfo (optional - auto-derived from app_id/name if not provided)
        provide_taskset_description: Function that returns taskset metadata
        provide_task_instances: Function that yields TaskInfo instances for given seeds
        rollout: Function that executes a rollout request and returns response
        dataset_registry: Optional registry for task datasets
        rubrics: Optional rubrics for evaluation
        proxy: Optional proxy config for vendor endpoints
        routers: Additional FastAPI routers to include
        middleware: Additional middleware to apply
        app_state: Mutable state dict accessible to routes
        require_api_key: Whether to require API key authentication
        ensure_localapi_auth: Whether to auto-ensure ENVIRONMENT_API_KEY on startup
        expose_debug_env: Whether to expose debug environment endpoint
        cors_origins: CORS allowed origins
        startup_hooks: Functions to call on app startup
        shutdown_hooks: Functions to call on app shutdown
    """

    app_id: str
    name: str
    description: str
    provide_taskset_description: TasksetDescriptor
    provide_task_instances: InstanceProvider
    rollout: RolloutExecutor
    base_task_info: TaskInfo | None = None  # Auto-derived from app_id/name if not provided
    dataset_registry: TaskDatasetRegistry | None = None
    rubrics: RubricBundle | None = field(default_factory=RubricBundle)
    proxy: ProxyConfig | None = None
    routers: Sequence[APIRouter] = field(default_factory=tuple)
    middleware: Sequence[Middleware] = field(default_factory=tuple)
    app_state: MutableMapping[str, Any] = field(default_factory=_default_app_state)
    require_api_key: bool = True
    ensure_localapi_auth: bool = True
    expose_debug_env: bool = True
    cors_origins: Sequence[str] | None = None
    startup_hooks: Sequence[Callable[[], None | Awaitable[None]]] = field(default_factory=tuple)
    shutdown_hooks: Sequence[Callable[[], None | Awaitable[None]]] = field(default_factory=tuple)

    def clone(self) -> TaskAppConfig:
        """Return a shallow copy safe to mutate when wiring the app."""

        return TaskAppConfig(
            app_id=self.app_id,
            name=self.name,
            description=self.description,
            provide_taskset_description=self.provide_taskset_description,
            provide_task_instances=self.provide_task_instances,
            rollout=self.rollout,
            base_task_info=self.base_task_info,  # May be None - auto-derived in create_task_app
            dataset_registry=self.dataset_registry,
            rubrics=self.rubrics or RubricBundle(),
            proxy=self.proxy,
            routers=tuple(self.routers),
            middleware=tuple(self.middleware),
            app_state=dict(self.app_state),
            require_api_key=self.require_api_key,
            ensure_localapi_auth=self.ensure_localapi_auth,
            expose_debug_env=self.expose_debug_env,
            cors_origins=tuple(self.cors_origins or ()),
            startup_hooks=tuple(self.startup_hooks),
            shutdown_hooks=tuple(self.shutdown_hooks),
        )


class LocalAPIConfig(TaskAppConfig):
    """Alias for TaskAppConfig with LocalAPI naming."""


def _maybe_await(result: Any) -> Awaitable[Any]:
    if inspect.isawaitable(result):
        return asyncio.ensure_future(result)
    loop = asyncio.get_event_loop()
    future: asyncio.Future[Any] = loop.create_future()
    future.set_result(result)
    return future


def _ensure_task_info(obj: Any) -> TaskInfo:
    if isinstance(obj, TaskInfo):
        return obj
    if isinstance(obj, MutableMapping):
        return TaskInfo.model_validate(obj)
    raise TypeError(
        f"Task instance provider must yield TaskInfo-compatible objects (got {type(obj)!r})"
    )


def _normalise_seeds(values: Sequence[int]) -> list[int]:
    seeds: list[int] = []
    for value in values:
        try:
            seeds.append(int(value))
        except Exception as exc:  # pragma: no cover - defensive
            raise ValueError(f"Seed values must be convertible to int (got {value!r})") from exc
    return seeds


def _build_proxy_routes(
    app: FastAPI, config: TaskAppConfig, auth_dependency: Callable[[Request], None]
) -> None:
    proxy = config.proxy
    if not proxy:
        return

    async def _call_vendor(
        url: str, payload: dict[str, Any], headers: dict[str, str]
    ) -> dict[str, Any]:
        async with httpx.AsyncClient(timeout=httpx.Timeout(600.0), follow_redirects=True) as client:
            response = await client.post(url, json=payload, headers=headers)
        data = (
            response.json()
            if response.headers.get("content-type", "").startswith("application/json")
            else {"raw": response.text}
        )
        if response.status_code >= 400:
            code = "vendor_error"
            if url.startswith("https://api.openai.com"):
                code = "openai_error"
            elif "groq" in url:
                code = "groq_error"
            raise http_exception(
                response.status_code,
                code,
                "Vendor proxy error",
                extra={"status": response.status_code, "body": data},
            )
        return data

    def _log_proxy(route: str, payload: dict[str, Any]) -> None:
        try:
            messages = payload.get("messages") if isinstance(payload, dict) else None
            msg_count = len(messages) if isinstance(messages, list) else 0
            tool_count = len(payload.get("tools") or []) if isinstance(payload, dict) else 0
            model = payload.get("model") if isinstance(payload, dict) else None
            print(
                f"[task:proxy:{route}] model={model} messages={msg_count} tools={tool_count}",
                flush=True,
            )
        except Exception:  # pragma: no cover - best effort logging
            pass

    system_hint = proxy.system_hint

    if proxy.enable_openai:

        @app.post("/proxy/v1/chat/completions", dependencies=[Depends(auth_dependency)])
        async def proxy_openai(body: dict[str, Any], request: Request) -> Any:  # type: ignore[no-redef]
            key = get_openai_key_or_503()
            model = body.get("model") if isinstance(body.get("model"), str) else None
            payload = prepare_for_openai(model, body)
            payload = inject_system_hint(payload, system_hint or "")
            _log_proxy("openai", payload)
            data = await _call_vendor(proxy.openai_url, payload, {"Authorization": f"Bearer {key}"})
            sanitized = synthesize_tool_call_if_missing(data)
            return to_jsonable(sanitized)

    if proxy.enable_groq:

        @app.post("/proxy/groq/v1/chat/completions", dependencies=[Depends(auth_dependency)])
        async def proxy_groq(body: dict[str, Any], request: Request) -> Any:  # type: ignore[no-redef]
            key = get_groq_key_or_503()
            model = body.get("model") if isinstance(body.get("model"), str) else None
            payload = prepare_for_groq(model, body)
            payload = inject_system_hint(payload, system_hint or "")
            _log_proxy("groq", payload)
            data = await _call_vendor(
                proxy.groq_url.rstrip("/"), payload, {"Authorization": f"Bearer {key}"}
            )
            sanitized = synthesize_tool_call_if_missing(data)
            return to_jsonable(sanitized)


def _auth_dependency_factory(config: TaskAppConfig) -> Callable[[Request], None]:
    def _dependency(request: Request) -> None:
        if not config.require_api_key:
            return
        require_api_key_dependency(request)

    return _dependency


def create_task_app(config: TaskAppConfig) -> FastAPI:
    """Create a FastAPI application from a TaskAppConfig.

    This function builds a complete FastAPI application that implements the Synth AI
    task app contract. The resulting app includes:
    - Health check endpoints (/health, /)
    - Task info endpoint (/task_info)
    - Rollout endpoint (/rollout)
    - Optional proxy endpoints for OpenAI/Groq
    - Optional dataset registry endpoints
    - CORS middleware (if configured)
    - Authentication middleware (if require_api_key=True)

    The app is ready to run with uvicorn or deploy to any ASGI-compatible server.

    Args:
        config: TaskAppConfig describing the task app

    Returns:
        FastAPI application instance ready to run

    Example:
        >>> from synth_ai.sdk.task.server import TaskAppConfig, create_task_app
        >>>
        >>> def build_config() -> TaskAppConfig:
        ...     return TaskAppConfig(
        ...         app_id="my_task",
        ...         name="My Task",
        ...         description="A task app",
        ...         provide_taskset_description=lambda: {"splits": ["train"]},
        ...         provide_task_instances=lambda seeds: [...],
        ...         rollout=lambda req, r: {...},
        ...         # base_task_info is optional - auto-derived from app_id/name
        ...     )
        >>>
        >>> app = create_task_app(build_config())
        >>> # Run with: uvicorn.run(app, host="0.0.0.0", port=8001)
    """
    ctx: dict[str, Any] = {"app_id": config.app_id, "name": config.name}
    log_info("create_task_app invoked", ctx=ctx)
    cfg = config.clone()
    cfg.rubrics = cfg.rubrics or RubricBundle()

    # Auto-derive base_task_info from app_id/name if not provided
    if cfg.base_task_info is None:
        cfg.base_task_info = TaskInfo(
            task={"id": cfg.app_id, "name": cfg.name},
            dataset={"id": cfg.app_id},
            inference={},
            limits={},
        )

    app = FastAPI(title=cfg.name, description=cfg.description)

    for key, value in cfg.app_state.items():
        setattr(app.state, key, value)

    if cfg.cors_origins is not None:
        app.add_middleware(
            CORSMiddleware,  # type: ignore[arg-type]
            allow_origins=list(cfg.cors_origins) or ["*"],
            allow_credentials=True,
            allow_methods=["*"],
            allow_headers=["*"],
        )

    # Note: additional middleware from cfg.middleware is currently disabled to avoid typing ambiguity.
    # for middleware in cfg.middleware:
    #     try:
    #         opts = getattr(middleware, "options", {})
    #     except Exception:
    #         opts = {}
    #     app.add_middleware(middleware.cls, **(opts if isinstance(opts, dict) else {}))

    for router in cfg.routers:
        try:
            app.include_router(router)
        except Exception:
            try:
                inner = getattr(router, "router", None)
                if inner is not None:
                    app.include_router(inner)
            except Exception:
                raise

    auth_dependency = _auth_dependency_factory(cfg)

    def _call_hook(hook: Callable[..., Any]) -> Awaitable[Any]:
        try:
            params = inspect.signature(hook).parameters  # type: ignore[arg-type]
        except (TypeError, ValueError):
            params = {}
        if params:
            return _maybe_await(hook(app))  # type: ignore[misc]
        return _maybe_await(hook())

    @asynccontextmanager
    async def lifespan(_: FastAPI):
        # Auto-ensure environment API key if enabled (defaults to True)
        # This handles minting, persisting, and uploading the key transparently
        if cfg.ensure_localapi_auth:
            from synth_ai.sdk.localapi.auth import ensure_localapi_auth

            ensure_localapi_auth()
        normalize_environment_api_key()
        normalize_vendor_keys()
        for hook in cfg.startup_hooks:
            await _call_hook(hook)
        try:
            yield
        finally:
            for hook in cfg.shutdown_hooks:
                await _call_hook(hook)

    app.router.lifespan_context = lifespan

    @app.get("/")
    async def root() -> Mapping[str, Any]:
        return to_jsonable({"status": "ok", "service": cfg.app_id})

    @app.head("/")
    async def root_head() -> Mapping[str, Any]:
        return to_jsonable({"status": "ok"})

    @app.get("/health", dependencies=[Depends(auth_dependency)])
    async def health(request: Request) -> Mapping[str, Any]:
        # If we got here, auth_dependency already verified the key exactly matches
        expected = normalize_environment_api_key()
        return to_jsonable(
            {
                "healthy": True,
                "auth": {
                    "required": True,
                    "expected_prefix": (expected[:6] + "...") if expected else "<unset>",
                },
            }
        )

    @app.post("/done", dependencies=[Depends(auth_dependency)])
    async def done() -> Mapping[str, Any]:
        # Coordination endpoint for tests and automation; indicates app is reachable
        return to_jsonable({"ok": True, "service": cfg.app_id})

    @app.get("/info", dependencies=[Depends(auth_dependency)])
    async def info() -> Mapping[str, Any]:
        dataset_meta = cfg.base_task_info.dataset
        rubrics: dict[str, Any] | None = None
        rubric_bundle = cfg.rubrics
        if rubric_bundle and (rubric_bundle.outcome or rubric_bundle.events):
            rubrics = {
                "outcome": rubric_bundle.outcome.model_dump() if rubric_bundle.outcome else None,
                "events": rubric_bundle.events.model_dump() if rubric_bundle.events else None,
            }
        payload = {
            "service": {
                "task": cfg.base_task_info.task,
                "version": cfg.base_task_info.task.version,
            },
            "dataset": dataset_meta,
            "rubrics": rubrics,
            "inference": cfg.base_task_info.inference,
            "limits": cfg.base_task_info.limits,
        }
        return to_jsonable(payload)

    @app.get("/task_info", dependencies=[Depends(auth_dependency)])
    async def task_info(
        seed: Sequence[int] | None = Query(default=None),
        seeds: Sequence[int] | None = Query(default=None),
    ) -> Any:
        all_seeds: list[int] = []
        if seed:
            all_seeds.extend(_normalise_seeds(seed))
        if seeds:
            all_seeds.extend(_normalise_seeds(seeds))

        if not all_seeds:
            descriptor_result = await _maybe_await(cfg.provide_taskset_description())
            return to_jsonable({"taskset": descriptor_result})

        instances = await _maybe_await(cfg.provide_task_instances(all_seeds))
        payload = [to_jsonable(_ensure_task_info(instance).model_dump()) for instance in instances]
        return payload[0] if len(payload) == 1 else payload

    @app.post("/rollout", dependencies=[Depends(auth_dependency)])
    async def rollout_endpoint(rollout_request: RolloutRequest, request: Request) -> Any:
        result = await _maybe_await(cfg.rollout(rollout_request, request))
        if isinstance(result, RolloutResponse):
            return to_jsonable(result.model_dump())
        if isinstance(result, Mapping):
            try:
                validated = RolloutResponse.model_validate(result)
            except Exception:
                return to_jsonable(result)
            return to_jsonable(validated.model_dump())
        raise TypeError("Rollout executor must return RolloutResponse or mapping")

    if cfg.expose_debug_env:

        @app.get("/debug/env", dependencies=[Depends(auth_dependency)])
        async def debug_env() -> Mapping[str, Any]:
            def _mask(value: str | None) -> str:
                if not value:
                    return ""
                return f"{value[:6]}…" if len(value) > 6 else value

            return to_jsonable(
                {
                    "has_ENVIRONMENT_API_KEY": bool(os.getenv("ENVIRONMENT_API_KEY")),
                    "OPENAI_API_KEY_prefix": _mask(os.getenv("OPENAI_API_KEY")),
                    "GROQ_API_KEY_prefix": _mask(os.getenv("GROQ_API_KEY")),
                }
            )

    _build_proxy_routes(app, cfg, auth_dependency)

    return app


def _load_env_files(env_files: Sequence[str]) -> list[str]:
    loaded: list[str] = []
    if not env_files:
        return loaded
    try:
        import dotenv
    except Exception:  # pragma: no cover - optional dep
        return loaded
    for path_str in env_files:
        path = Path(path_str)
        if not path.is_file():
            continue
        dotenv.load_dotenv(path, override=False)
        loaded.append(str(path))
    return loaded


def run_task_app(
    config_factory: Callable[[], TaskAppConfig],
    *,
    host: str = "0.0.0.0",
    port: int = 8001,
    reload: bool = False,
    env_files: Sequence[str] = (),
) -> None:
    """Run a task app locally with uvicorn.

    This is a convenience function that creates a task app from a config factory
    and runs it with uvicorn. It handles environment file loading and service
    registration for tunnel management.

    Args:
        config_factory: Function that returns a TaskAppConfig
        host: Host to bind to (default: "0.0.0.0")
        port: Port to listen on (default: 8001)
        reload: Enable auto-reload on file changes (default: False)
        env_files: Paths to .env files to load (default: empty)

    Example:
        >>> from synth_ai.sdk.task.server import run_task_app
        >>> from my_task_app import build_config
        >>>
        >>> run_task_app(
        ...     build_config,
        ...     host="127.0.0.1",
        ...     port=8114,
        ...     reload=True,
        ...     env_files=[".env.local"],
        ... )

    Raises:
        RuntimeError: If uvicorn is not installed
        TypeError: If config_factory doesn't return TaskAppConfig
    """
    ctx: dict[str, Any] = {"host": host, "port": port, "reload": reload}
    log_info("run_task_app invoked", ctx=ctx)

    loaded_files = _load_env_files(env_files)
    if loaded_files:
        print(f"[task:server] Loaded environment from: {', '.join(loaded_files)}", flush=True)

    config = config_factory()
    # Defensive: ensure the factory produced a valid TaskAppConfig to avoid
    # confusing attribute errors later in the boot sequence.
    if not isinstance(config, TaskAppConfig):  # type: ignore[arg-type]
        raise TypeError(
            f"Task app config_factory must return TaskAppConfig, got {type(config).__name__}"
        )
    app = create_task_app(config)

    try:
        import uvicorn
    except ImportError as exc:  # pragma: no cover - uvicorn optional
        raise RuntimeError("uvicorn must be installed to run the task app locally") from exc

    print(f"[task:server] Starting '{config.app_id}' on {host}:{port}", flush=True)
    
    # Record local service before starting
    try:
        import os

        from synth_ai.cli.lib.tunnel_records import record_service
        local_url = f"http://{host if host not in ('0.0.0.0', '::') else '127.0.0.1'}:{port}"
        # Try to get current process PID
        pid: int | None = os.getpid()
        record_service(
            url=local_url,
            port=port,
            service_type="local",
            local_host=host if host not in ('0.0.0.0', '::') else '127.0.0.1',
            app_id=config.app_id,
            pid=pid,
        )
    except Exception:
        pass  # Fail silently - records are optional
    
    try:
        uvicorn.run(app, host=host, port=port, reload=reload)
    finally:
        # Clean up record when server exits
        try:
            from synth_ai.cli.lib.tunnel_records import remove_service_record
            remove_service_record(port)
        except Exception:
            pass


def run_server_background(
    app: Any,
    port: int,
    host: str = "0.0.0.0",
) -> threading.Thread:
    """Start uvicorn server in a background daemon thread.

    For manual control over task app lifecycle. If you want automatic
    tunnel management, use InProcessTaskApp instead.

    Args:
        app: ASGI/FastAPI application
        port: Port to bind
        host: Host to bind (default 0.0.0.0 for tunnel access)

    Returns:
        Daemon thread running the server (stops when main process exits)

    Example:
        from synth_ai.sdk.task import run_server_background
        from synth_ai.sdk.tunnels import wait_for_health_check

        thread = run_server_background(my_app, port=8001)
        await wait_for_health_check("localhost", 8001, api_key)
    """
    import asyncio
    import threading

    import uvicorn

    def serve() -> None:
        config = uvicorn.Config(app, host=host, port=port, log_level="warning")
        server = uvicorn.Server(config)
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        loop.run_until_complete(server.serve())

    thread = threading.Thread(target=serve, daemon=True, name=f"uvicorn-{port}")
    thread.start()
    return thread
