import functools
import inspect
import json
import logging
import random
import re
import string
import threading
import uuid
from datetime import datetime, timedelta, timezone
from typing import Any, Callable, Dict, List, Optional, Type

import openai
from fastapi import APIRouter, Depends, HTTPException, Query
from fastapi.responses import JSONResponse
from fastapi_utils.cbv import cbv
from fastapi_utils.inferring_router import InferringRouter

from api.api_schemas import (
    ApplicationBlocksResponse,
    ApplicationConfiguration,
    ApplicationCreate,
    ApplicationCreateResponse,
    ApplicationID,
    ApplicationListResponse,
    ApplicationRun,
    ApplicationRunResponse,
    ApplicationTypesResponse,
    ApplicationUsageResponse,
    AppWithLatestCall,
    BlockInfo,
    DebugInfo,
    InteractionDebugResponse,
    ItemDeleteResponse,
    ItemUpdateResponse,
    Metadata,
    Parameter,
    TypeInfo,
    UsageInfo,
)
from api.dependencies import get_logger
from database import (
    create_application,
    create_interaction,
    delete_application,
    delete_interactions,
    filter_applications,
    get_application,
    get_interaction,
    get_interactions,
    ping_db,
    update_application,
    update_interaction,
)
from exceptions import (
    ApplicationInputTypeMismatch,
    ApplicationNotFound,
    InteractionNotFound,
    NodeConstructError,
    register_exception_handlers,
)
from models import Application, Interaction
from resolver.resolver import Resolver
from scheduler.graph import Edge, Graph

router = InferringRouter()

name_resolver = Resolver("./names.yaml")


def construct_node(config: dict) -> Any:
    if not isinstance(config, dict):
        name = config.name
    else:
        name = config["name"]
    cls = name_resolver.lookup(name)
    if cls is None:
        raise NodeConstructError(f"name {name} not found")
    elif name_resolver.is_abstract(cls):
        raise NodeConstructError(
            f"{name} is an abstract type and can NOT be constructed"
        )
    properties = {}
    if not isinstance(config, dict):
        slots = config.slots
        if slots is None:
            slots = {}
    else:
        slots = config.get("slots", {})
    for p, v in slots.items():
        if isinstance(v, dict):
            properties[p] = construct_node(v)
        elif isinstance(v, list):
            properties[p] = [construct_node(x) if isinstance(x, dict) else x for x in v]
        else:
            properties[p] = v

    try:
        return cls(**properties)
    except Exception as e:
        raise NodeConstructError(f"construct {name} failed: {str(e)}") from e


def resolve_params(params: Dict[str, inspect.Parameter]):
    ss = []
    for sname, param in params.items():
        x = {
            "name": sname,
            "class": name_resolver.relookup(param.annotation),
        }
        p = Parameter(
            name=sname,
            class_name=name_resolver.relookup(param.annotation),
        )
        if param.default != inspect.Parameter.empty:
            p.default = param.default
        if param.kind == param.VAR_KEYWORD:
            p.is_variable_keyword = True
        ss.append(p)
    return ss


def initialize_app(configuration: dict, skip_validation: bool = False):
    edges = []
    for edge in configuration["edges"]:
        edges.append(
            Edge(
                source=edge.get("src_block"),
                sink=edge.get("dst_block"),
                port=edge.get("dst_port"),
                case=edge.get("case"),
            )
        )
    nodes = {}
    for node in configuration["nodes"]:
        nodes[node.get("id")] = construct_node(node)
    return Graph(nodes, edges, skip_validation)


def add_random_suffix(input_string):
    suffix = "".join(
        random.choices(
            string.ascii_lowercase + string.digits + string.ascii_uppercase, k=8
        )
    )
    return f"{input_string}_{suffix}"


def convert_to_dict(obj):
    if isinstance(obj, dict):
        return {
            key: convert_to_dict(value)
            for key, value in obj.items()
            if value is not None
        }
    elif isinstance(obj, list):
        return [convert_to_dict(item) for item in obj]
    elif hasattr(obj, "__dict__"):
        return convert_to_dict(obj.__dict__)
    else:
        return obj


def remove_suffix_version(input_string):
    pattern = r"_v\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}\.\d{6}$"
    match = re.search(pattern, input_string)
    if match:
        return match.group(0), input_string[: -len(match.group(0))]
    else:
        return None, input_string


def json_serialize(instance):
    # default supported classes: https://github.com/python/cpython/blob/main/Lib/json/encoder.py#L77
    default_serializable = (dict, list, tuple, str, int, float, bool)
    if instance is None:
        return None
    elif isinstance(instance, default_serializable):
        return instance
    elif hasattr(instance, "json_serialize"):
        return instance.json_serialize()
    else:
        raise NotImplementedError(
            f"Class {type(instance)} does not implement json_serialize()"
        )


def json_serialize_dict(input_dict):
    for key, value in input_dict.items():
        input_dict[key] = json_serialize(value)
    return input_dict


class AsyncExceptionHandler:
    def __init__(self):
        self.exception_handlers = {}

    def add_exception_handler(self, exc_class: Type[Exception], handler: Callable):
        self.exception_handlers[exc_class] = handler

    def exception_handler(self, exc_class: Type[Exception]):
        return functools.partial(self.add_exception_handler, exc_class)

    def render(self, exc) -> JSONResponse:
        exc_class = type(exc)

        if exc_class in self.exception_handlers:
            return self.exception_handlers[exc_class](None, exc)
        return self.exception_handlers[Exception](None, exc)


@cbv(router)
class ApplicationView:
    @router.post("/applications")
    def new_application(
        self,
        application: ApplicationCreate,
        deploy: bool = Query(False, description="deploy"),
        logger: logging.Logger = Depends(get_logger),
    ) -> ApplicationCreateResponse:
        config = convert_to_dict(application.configuration)
        metadata = config.get("metadata")
        initialize_app(config)
        _id = str(uuid.uuid4())

        # unique_name = add_random_suffix(metadata.get("name"))

        created_at = datetime.utcnow()

        config["metadata"]["created_at"] = int(created_at.timestamp())
        config["metadata"]["updated_at"] = int(created_at.timestamp())

        unique_name = metadata.get("name") + "_v" + created_at.isoformat()

        app = Application(
            id=_id,
            name=unique_name,
            owner=metadata.get("owner"),
            created_at=created_at,
            updated_at=created_at,
            version=metadata.get("version"),
            description=metadata.get("description"),
            deploy=deploy,
            custom_fields=metadata.get("custom_fields"),
            configuration=config,
        )
        create_application(app)
        logger.info({"event": "application.create", "application_id": _id})
        return ApplicationCreateResponse(id=_id)

    @router.post("/applications/{application_id}")
    def run_application(
        self,
        app: ApplicationRun,
        application_id: str,
        return_trace: bool = Query(False, description="return trace data"),
        logger: logging.Logger = Depends(get_logger),
    ) -> ApplicationRunResponse:
        app_config = get_application(application_id)
        if app_config is None:
            raise ApplicationNotFound(application_id)
        app_instance = initialize_app(app_config.configuration, skip_validation=True)

        if not isinstance(app.input, app_instance.input_type()):
            raise ApplicationInputTypeMismatch(
                app_instance.input_type(), type(app.input)
            )
        output = app_instance.run(app.input)
        logs = app_instance.data
        logs["id"] = str(uuid.uuid4())
        logger.info({**logs})

        interaction = Interaction(
            id=logs["id"],
            created_at=datetime.utcnow(),
            application_id=application_id,
            data=json_serialize_dict(app_instance.data),
            usage=app_instance.usage,
            error=None,
        )
        create_interaction(interaction)

        trace = None
        if return_trace:
            trace = json_serialize_dict(app_instance.data)
        if app_instance.usage:
            usage = []
            for n, usage_data in app_instance.usage.items():
                usage.append(
                    UsageInfo(
                        block=n,
                        prompt_tokens=usage_data["prompt_tokens"],
                        completion_tokens=usage_data["completion_tokens"],
                        total_cost=usage_data["total_cost"],
                    )
                )
        else:
            usage = None
        return ApplicationRunResponse(
            id=logs["id"], output=output, trace=trace, usage=usage
        )

    @router.post("/applications/{application_id}/async")
    def run_application_async(
        self,
        app: ApplicationRun,
        application_id: str,
        return_trace: bool = Query(False, description="return trace data"),
        logger: logging.Logger = Depends(get_logger),
    ) -> ApplicationRunResponse:
        app_config = get_application(application_id)
        if app_config is None:
            raise ApplicationNotFound(application_id)
        app_instance = initialize_app(app_config.configuration, skip_validation=True)

        if not isinstance(app.input, app_instance.input_type()):
            raise ApplicationInputTypeMismatch(
                app_instance.input_type(), type(app.input)
            )

        interaction_id = str(uuid.uuid4())

        interaction = Interaction(
            id=interaction_id,
            created_at=datetime.utcnow(),
            application_id=application_id,
            data={},
            usage=None,
            error=None,
        )
        create_interaction(interaction)

        def async_run():
            h = AsyncExceptionHandler()
            register_exception_handlers(h)
            try:
                output = app_instance.run(app.input)
                logs = app_instance.data
                logs["id"] = interaction_id
                logger.info({**logs})

                interaction = Interaction(
                    id=logs["id"],
                    created_at=datetime.utcnow(),
                    application_id=application_id,
                    data=json_serialize_dict(app_instance.data),
                    usage=app_instance.usage,
                    error=None,
                )
                update_interaction(interaction)
            except Exception as e:
                r = h.render(e)
                interaction = Interaction(
                    id=interaction_id,
                    created_at=datetime.utcnow(),
                    application_id=application_id,
                    usage=app_instance.usage,
                    data=None,
                    error={
                        "status_code": r.status_code,
                        "content": json.loads(r.body),
                    },
                )
                update_interaction(interaction)

        threading.Thread(target=async_run).start()

        return ApplicationRunResponse(
            id=interaction_id, output=None, trace=None, usage=None
        )

    @router.get("/applications/{application_id}/debug/{interaction_id}")
    def debug_interaction(
        self,
        interaction_id: str,
        application_id: str,
        block: Optional[List[str]] = Query(None),
    ) -> InteractionDebugResponse:
        interaction = get_interaction(interaction_id)
        if interaction is None:
            raise InteractionNotFound(interaction_id)
        if interaction.error is not None:
            return JSONResponse(**interaction.error)
        logs = interaction.data
        result = []
        if block is None:
            block = logs.keys()
        for b in block:
            if b not in logs:
                result.append(DebugInfo(block=b, output="Not found"))
            else:
                result.append(DebugInfo(block=b, output=str(logs[b])))
        if interaction.usage:
            usage = []
            for n, usage_data in interaction.usage.items():
                usage.append(
                    UsageInfo(
                        block=n,
                        prompt_tokens=usage_data["prompt_tokens"],
                        completion_tokens=usage_data["completion_tokens"],
                        total_cost=usage_data["total_cost"],
                    )
                )
        else:
            usage = None
        return InteractionDebugResponse(
            application=application_id,
            interaction=interaction_id,
            debug=result,
            usage=usage,
        )

    @router.get("/applications/types")
    def types(self) -> ApplicationTypesResponse:
        names = name_resolver.names()
        types = []
        for name in names:
            if (
                name_resolver.lookup(name, "category") == "type"
                or name_resolver.lookup(name, "category") == "builtin"
            ):
                ti = TypeInfo(
                    name=name,
                    alias=name_resolver.lookup(name, "alias"),
                    candidates=name_resolver.candidates(name),
                    slots=None,
                )
                if (not name_resolver.is_abstract(name_resolver.lookup(name))) and (
                    name_resolver.lookup(name, "category") == "type"
                ):
                    ti.slots = resolve_params(name_resolver.slots(name))
                types.append(ti)

        return ApplicationTypesResponse(types=types)

    @router.get("/applications/blocks")
    def blocks(self) -> ApplicationBlocksResponse:
        names = name_resolver.names()
        blocks = []
        for name in names:
            if name_resolver.lookup(name, "category") == "block":
                blocks.append(
                    BlockInfo(
                        name=name,
                        alias=name_resolver.lookup(name, "alias"),
                        dir=name_resolver.lookup(name, "dir"),
                        slots=resolve_params(name_resolver.slots(name)),
                        inports=resolve_params(name_resolver.inports(name)),
                        outports=[
                            Parameter(
                                name="output",
                                class_name=name_resolver.relookup(
                                    name_resolver.outport(name)
                                ),
                            )
                        ],
                    )
                )

        return ApplicationBlocksResponse(blocks=blocks)

    @router.get("/applications/{application_id}")
    def query_app(self, application_id: str) -> ApplicationConfiguration:
        app = get_application(application_id)
        if app is None:
            raise ApplicationNotFound(application_id)
        return ApplicationConfiguration(**app.configuration)

    @router.delete("/applications/{application_id}")
    def del_app(self, application_id: str) -> ItemDeleteResponse:
        app = delete_application(application_id)
        delete_interactions(application_id)
        return ItemDeleteResponse(
            success=app,
            message=f"Application {application_id} and its interactions are deleted.",
        )

    @router.put("/applications/{application_id}")
    def update_app_meta(
        self, application_id: str, metadata: Metadata
    ) -> ItemUpdateResponse:
        """
        name: Without _v20230101120000000000
        owner:
        created_at: uneditable
        updated_at: uneditable

        """
        old_app = get_application(application_id)
        if old_app is None:
            raise ApplicationNotFound(application_id)
        config = old_app.configuration
        updated_at = datetime.utcnow()
        meta = config["metadata"]
        _, meta["name"] = (
            remove_suffix_version(metadata.name)
            if metadata.name != "Untitled_Application"
            else (None, meta["name"])
        )
        meta["owner"] = (
            metadata.owner if metadata.owner != "anonymous" else meta["owner"]
        )
        meta["updated_at"] = int(updated_at.timestamp())
        meta["deploy"] = metadata.deploy
        meta["version"] = (
            metadata.version if metadata.version is not None else meta.get("version")
        )
        meta["description"] = (
            metadata.description
            if metadata.description is not None
            else meta.get("description")
        )
        meta["custom_fields"] = (
            metadata.custom_fields
            if metadata.custom_fields != {}
            else meta["custom_fields"]
        )

        config["metadata"] = meta

        app = Application(
            id=application_id,
            name=meta["name"] + "_v" + updated_at.isoformat(),
            owner=meta["owner"],
            created_at=old_app.created_at,
            updated_at=updated_at,
            version=meta["version"],
            description=meta["description"],
            deploy=meta["deploy"],
            custom_fields=meta["custom_fields"],
            configuration=config,
        )
        update_application(app)
        return ItemUpdateResponse(
            success=True, message=f"Application {application_id}'s metadata updated."
        )

    @router.get("/applications")
    def filter_app(
        self,
        offset: int = Query(0, description="offset"),
        limit: int = Query(10, description="limit"),
        sort_by: str = Query("updated_at", description="sort by"),
        descend: bool = Query(True, description="descend"),
        owner: Optional[str] = Query(None, description="filter by owner"),
        name: Optional[str] = Query(None, description="filter by name"),
        version: Optional[str] = Query(None, description="filter by version"),
        description: Optional[str] = Query(None, description="filter by description"),
        deploy: Optional[bool] = Query(None, description="filter by deploy"),
    ) -> ApplicationListResponse:
        apps, total = filter_applications(
            offset,
            limit,
            sort_by,
            descend,
            owner=owner,
            name=name,
            version=version,
            description=description,
            deploy=deploy,
        )

        app_list_w_latest_call = [
            AppWithLatestCall(
                application_id=app["id"],
                name=app["name"],
                owner=app["owner"],
                created_at=int(app["created_at"].timestamp()),
                description=app["description"],
                deploy=app["deploy"],
                version=app["version"],
                custom_fields=app["custom_fields"],
                latest_called_at=app["latest_called_at"],
            )
            for app in apps
        ]
        return ApplicationListResponse(
            applications=app_list_w_latest_call,
            total=total,
        )

    @router.get("/applications/{application_id}/usage")
    def get_app_usage(
        self,
        application_id: str,
        _from: str = Query(None, alias="from", description="from date like 2022-01-01"),
        _to: str = Query(None, alias="to", description="to date like 2022-01-02"),
    ):
        _from = (
            datetime.strptime(_from, "%Y-%m-%d")
            if _from is not None
            else datetime(1970, 1, 1)
        )
        _from = _from.replace(tzinfo=timezone.utc)
        _to = (
            datetime.strptime(_to, "%Y-%m-%d") if _to is not None else datetime.utcnow()
        )
        _to = _to.replace(tzinfo=timezone.utc) + timedelta(days=1, seconds=-1)
        interactions = get_interactions(application_id, _from, _to)
        if interactions is None:
            return ApplicationUsageResponse(application_id, 0, 0, 0, 0, 0.0)
        (
            total_cost,
            prompt_tokens,
            completion_tokens,
            successful_llm_calls,
            successful_interactions,
        ) = (0, 0, 0, 0, 0)
        for interaction in interactions:
            if not interaction.error:
                if interaction.usage and len(interaction.usage) > 0:
                    for _, usage_data in interaction.usage.items():
                        prompt_tokens += usage_data["prompt_tokens"]
                        completion_tokens += usage_data["completion_tokens"]
                        total_cost += usage_data["total_cost"]
                        successful_llm_calls += 1
                successful_interactions += 1
        return ApplicationUsageResponse(
            id=application_id,
            successful_interactions=successful_interactions,
            successful_llm_calls=successful_llm_calls,
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
            total_cost=total_cost,
        )


@router.get("/ping")
def ping():
    ping_db()
    return {"message": "pong"}