{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}

module API.Docs.Generate.Python where

import API.Docs.Commands
import API.Docs.Events
import API.Docs.Generate
import API.Docs.Responses
import API.Docs.Syntax
import API.Docs.Syntax.Types
import API.Docs.Types
import API.TypeInfo
import Data.Char (isAlphaNum, toUpper)
import qualified Data.List.NonEmpty as L
import Data.Text (Text)
import qualified Data.Text as T

commandsCodeFile :: FilePath
commandsCodeFile = "./packages/simplex-chat-python/src/simplex_chat/types/_commands.py"

responsesCodeFile :: FilePath
responsesCodeFile = "./packages/simplex-chat-python/src/simplex_chat/types/_responses.py"

eventsCodeFile :: FilePath
eventsCodeFile = "./packages/simplex-chat-python/src/simplex_chat/types/_events.py"

typesCodeFile :: FilePath
typesCodeFile = "./packages/simplex-chat-python/src/simplex_chat/types/_types.py"

-- | Replace dashes with underscores so Python identifiers stay valid.
pyIdent :: String -> Text
pyIdent = T.replace "-" "_" . T.pack

-- | Python class name for a union member tag.
pyConstrName :: String -> Text
pyConstrName = pyIdent . fstToUpper

commandsCodeText :: Text
commandsCodeText =
  ("# API Commands\n# " <> autoGenerated <> "\n")
    <> "from __future__ import annotations\n"
    <> "import json\n"
    <> "from typing import NotRequired, TypedDict\n"
    <> "from . import _types as T\n"
    <> "from . import _responses as CR\n"
    <> foldMap commandCatCode chatCommandsDocs
  where
    commandCatCode CCCategory {categoryName, categoryDescr, commands} =
      (T.pack $ "\n# " <> categoryName <> "\n# " <> categoryDescr <> "\n")
        <> foldMap commandCode commands
      where
        commandCode CCDoc {commandType = ATUnionMember tag params, commandDescr, syntax, responses, network} =
          ("\n# " <> commandDescr <> "\n")
            <> ("# Network usage: " <> networkUsage network <> ".\n")
            <> classDef
            <> (if syntax == "" then "" else cmdStringFunc)
            <> respAliasLine
          where
            constrName = T.pack $ fstToUpper tag
            classDef =
              ("class " <> constrName <> "(TypedDict):\n")
                <> bodyOrPass (fieldsCodePy "    " "T." params)
                <> "\n"
            cmdStringFunc =
              ("\ndef " <> constrName <> "_cmd_string(self: " <> constrName <> ") -> str:\n")
                <> "    return " <> pySelfSyntaxText "T." (fstToUpper tag, params) syntax <> "\n"
            respAliasLine =
              "\n" <> constrName <> "_Response = " <> respUnion <> "\n"
            respUnion = unionAliasRhs "" (responseRef . responseType) responses
            responseRef (ATUnionMember rtag _) = "CR." <> pyConstrName rtag

responsesCodeText :: Text
responsesCodeText =
  ("# API Responses\n# " <> autoGenerated <> "\n")
    <> pythonImports
    <> unionTypeCodePy moduleMember "T." "ChatResponse" chatRespConstrs
  where
    chatRespConstrs = L.fromList $ map responseType chatResponsesDocs

eventsCodeText :: Text
eventsCodeText =
  ("# API Events\n# " <> autoGenerated <> "\n")
    <> "from __future__ import annotations\n"
    <> "from collections.abc import Awaitable, Callable\n"
    <> "from typing import Literal, NotRequired, Protocol, TypedDict, overload\n"
    <> "from . import _types as T\n"
    <> unionTypeCodePy moduleMember "T." "ChatEvent" chatEventConstrs
    <> onEventProtocolCode chatEventConstrs
  where
    chatEventConstrs = L.fromList $ concatMap catEvents chatEventsDocs
    catEvents CECategory {mainEvents, otherEvents} = map eventType $ mainEvents ++ otherEvents

-- | Render the `OnEventDecorator` Protocol — one `__call__` overload per
-- event tag, narrowing the handler's event parameter from the unnarrowed
-- `ChatEvent` union to the specific tagged TypedDict. Plus a fallback
-- overload for `event: str` that keeps the unnarrowed shape so non-literal
-- tags don't trigger a type error.
--
-- `Client.on_event` is typed as a `OnEventDecorator` (via a property) so
-- callers get per-tag narrowing without per-tag handwritten overloads
-- in client.py.
onEventProtocolCode :: L.NonEmpty ATUnionMember -> Text
onEventProtocolCode members =
  "\n\nclass OnEventDecorator(Protocol):\n"
    <> "    \"\"\"Per-tag narrowing protocol for ``Client.on_event``.\n"
    <> "\n"
    <> "    ``@client.on_event(\"contactConnected\")`` types the handler's\n"
    <> "    ``evt`` parameter as :class:`ContactConnected` rather than the\n"
    <> "    unnarrowed :data:`ChatEvent` union.\n"
    <> "    \"\"\"\n"
    <> foldMap overloadCode (L.toList members)
    <> "\n    @overload\n"
    <> "    def __call__(self, event: str, /) -> Callable[\n"
    <> "        [Callable[[\"ChatEvent\"], Awaitable[None]]],\n"
    <> "        Callable[[\"ChatEvent\"], Awaitable[None]],\n"
    <> "    ]: ...\n"
  where
    overloadCode (ATUnionMember tag _) =
      "\n    @overload\n"
        <> "    def __call__(self, event: Literal[\"" <> T.pack tag <> "\"], /) -> Callable[\n"
        <> "        [Callable[[\"" <> pyConstrName tag <> "\"], Awaitable[None]]],\n"
        <> "        Callable[[\"" <> pyConstrName tag <> "\"], Awaitable[None]],\n"
        <> "    ]: ...\n"

typesCodeText :: Text
typesCodeText =
  ("# API Types\n# " <> autoGenerated <> "\n")
    <> "from __future__ import annotations\n"
    <> "from typing import Literal, NotRequired, TypedDict\n"
    <> foldMap typeCode chatTypesDocs
  where
    typeCode ctd@CTDoc {typeDef = APITypeDef {typeName' = name, typeDef}, typeDescr} =
      (if T.null typeDescr then "" else "\n# " <> typeDescr <> "\n")
        <> typeDefCode
        <> typeCmdStringCode ctd
      where
        name' = T.pack name
        enumValue m = case name of
          "ConnectionMode" -> map toUpper m
          "FileProtocol" -> map toUpper m
          _ -> m
        typeDefCode = case typeDef of
          ATDRecord fields ->
            ("\nclass " <> name' <> "(TypedDict):\n")
              <> bodyOrPass (fieldsCodePy "    " "" fields)
          ATDEnum cs ->
            "\n" <> name' <> " = Literal["
              <> T.intercalate ", " (map (\m -> "\"" <> T.pack (enumValue m) <> "\"") $ L.toList cs)
              <> "]\n"
          ATDUnion cs -> unionTypeCodePy typeMember "" name cs

-- | For types with non-empty `typeSyntax`, emit a top-level
-- `<TypeName>_cmd_string(self: <TypeName>) -> str` helper that mirrors the
-- Choice/Param expression. Records access fields via `self['<name>']`;
-- enums and unions dispatch on `self` (a literal string) or `self['type']`
-- respectively. Required so generated `_commands.py` produces valid CLI
-- syntax for ChatRef/ChatType/ChatDeleteMode/GroupChatScope/PaginationByTime
-- params instead of stringifying the wire dict.
typeCmdStringCode :: CTDoc -> Text
typeCmdStringCode CTDoc {typeDef = td@APITypeDef {typeName' = name, typeDef}, typeSyntax}
  | typeSyntax == "" = ""
  | otherwise =
      "\n\ndef " <> T.pack name <> "_cmd_string(self: " <> T.pack name <> ") -> str:\n"
        <> "    return " <> body <> ignore <> "\n"
  where
    body = pyTypeSyntaxText "" (name, fields) typeSyntax
    -- Unions and enums use self/self['type'] to dispatch. Pyright cannot
    -- narrow TypedDict access by string-literal key, so suppress per-branch
    -- complaints with one ignore on the return.
    ignore = case typeDef of
      ATDUnion _ -> "  # type: ignore[typeddict-item]"
      _ -> ""
    -- typeFields mirrors TS funcCode: include `self` so Choice "self"
    -- resolves; for unions add `type` and flatten member fields.
    self = APIRecordField "self" (ATDef td)
    fields = case typeDef of
      ATDRecord fs -> fs
      ATDUnion ms ->
        self : APIRecordField "type" tagType : concatMap (\(ATUnionMember _ fs) -> fs) (L.toList ms)
        where
          tagType = ATDef $ APITypeDef (name <> ".type") $ ATDEnum tags
          tags = L.map (\(ATUnionMember tag _) -> tag) ms
      ATDEnum _ -> [self]

-- | Like `pySelfSyntaxText` but excludes `self` from the param-rewrite list
-- so `self == 'tag'` (enum dispatch) and `self['type']` (union dispatch)
-- survive verbatim. Used only for type-level cmd_string functions inside
-- @_types.py@, where peer type cmd_string calls don't need a namespace.
pyTypeSyntaxText :: String -> TypeAndFields -> Expr -> Text
pyTypeSyntaxText typeNamespace r expr =
  rewriteParams accessors (pySyntaxText typeNamespace r expr)
  where
    accessors = filter ((/= "self") . fst) (paramAccessors r)

-- | Member class name within the multi-type @_types.py@ module: prefix the
-- tag with the union type name so members from different unions don't
-- collide.
typeMember :: String -> String -> Text
typeMember typeName tag = T.pack typeName <> "_" <> pyIdent tag

-- | Member class name within a single-union module (responses/events): just
-- the PascalCase tag, so commands can reference them as @CR.<Tag>@.
moduleMember :: String -> String -> Text
moduleMember _ tag = pyConstrName tag

-- | Common imports for the responses/events modules.
pythonImports :: Text
pythonImports =
  "from __future__ import annotations\n"
    <> "from typing import Literal, NotRequired, TypedDict\n"
    <> "from . import _types as T\n"

-- | Render a tagged-union type: one TypedDict per member, plus union alias
-- and `<Name>_Tag` Literal alias. The member class names are produced by
-- @memberName@ given the union type name and the member tag.
unionTypeCodePy ::
  (String -> String -> Text) ->
  Text ->
  String ->
  L.NonEmpty ATUnionMember ->
  Text
unionTypeCodePy memberName typesNamespace name cs =
  foldMap memberClass (L.toList cs)
    <> "\n" <> name' <> " = " <> unionAliasRhs name' constrTypeRef (L.toList cs)
    <> "\n" <> name' <> "_Tag = Literal[" <> tagLiterals <> "]\n"
  where
    name' = T.pack name
    constrTypeRef (ATUnionMember tag _) = memberName name tag
    tagLiterals = T.intercalate ", " $ map (\(ATUnionMember tag _) -> "\"" <> T.pack tag <> "\"") $ L.toList cs
    memberClass (ATUnionMember tag fields) =
      ("\nclass " <> memberName name tag <> "(TypedDict):\n")
        <> ("    type: Literal[\"" <> T.pack tag <> "\"]\n")
        <> fieldsCodePy "    " typesNamespace fields

-- | Render the right-hand side of a union alias: either inline (one line) or
-- multi-line wrapped in parentheses with `|` separators between alternatives.
unionAliasRhs :: Text -> (a -> Text) -> [a] -> Text
unionAliasRhs lhs constr cs
  | T.length (lhs <> " = " <> oneLine) <= 100 = oneLine <> "\n"
  | otherwise = "(\n" <> T.intercalate "\n" (map ("    " <>) lines') <> "\n)\n"
  where
    oneLine = T.intercalate " | " cs'
    lines' = case cs' of
      [] -> []
      (h : t) -> h : map ("| " <>) t
    cs' = map constr cs

-- | Emit a body of `pass` if there are no fields, otherwise the rendered
-- fields as-is.
bodyOrPass :: Text -> Text
bodyOrPass body
  | T.null body = "    pass\n"
  | otherwise = body

-- | Render record fields for a TypedDict body. Each field becomes
-- `<indent><name>: <type>[  # <comment>]`. Optional fields wrap the type in
-- `NotRequired[...]`.
fieldsCodePy :: Text -> Text -> [APIRecordField] -> Text
fieldsCodePy indent namespace = foldMap render
  where
    render (APIRecordField name t) =
      indent <> T.pack name <> ": " <> wrapOptional t (typeText t) <> typeComment t <> "\n"
    wrapOptional t inner = case t of
      ATOptional _ -> "NotRequired[" <> inner <> "]"
      _ -> inner
    typeText = \case
      ATPrim (PT t) -> primName t
      ATDef (APITypeDef t _) -> quoted (namespace <> T.pack t)
      ATRef t -> quoted (namespace <> T.pack t)
      ATOptional t -> typeText t
      ATArray {elemType} -> "list[" <> typeText elemType <> "]"
      ATMap (PT k) v -> "dict[" <> primName k <> ", " <> typeText v <> "]"
    primName = \case
      TBool -> "bool"
      TString -> "str"
      TInt -> "int"
      TInt64 -> "int"
      TWord32 -> "int"
      TDouble -> "float"
      TJSONObject -> "dict[str, object]"
      TUTCTime -> "str"
      t -> T.pack t
    quoted s = "\"" <> s <> "\""
    typeComment t = let c = typeComment' t in if T.null c then "" else "  # " <> c
    typeComment' = \case
      ATPrim (PT t) -> typeComment_ t
      ATOptional inner -> typeComment' inner
      ATArray {elemType, nonEmpty}
        | nonEmpty -> if T.null c then "non-empty" else c <> ", non-empty"
        | otherwise -> c
        where
          c = typeComment' elemType
      ATMap (PT k) v ->
        let kc = typeComment_ k
            vc = typeComment' v
            tc t c = if T.null c then t else c
         in if T.null kc && T.null vc then "" else tc (primName k) kc <> " : " <> tc (typeText v) vc
      _ -> ""
    typeComment_ = \case
      TInt -> "int"
      TInt64 -> "int64"
      TWord32 -> "word32"
      TDouble -> "double"
      TUTCTime -> "ISO-8601 timestamp"
      _ -> ""

-- | Wrap `pySyntaxText` so each parameter access uses `self['<name>']`. The
-- output of `pySyntaxText` references params as bare Python identifiers
-- (e.g. `str(userId)`); we rewrite those identifiers — but only outside
-- string literals — into TypedDict subscript accesses. The
-- @typeNamespace@ is prepended to any `<TypeName>_cmd_string(...)` calls
-- emitted for params whose type has its own syntax (e.g. @"T."@ from
-- @_commands.py@, or @""@ from within @_types.py@).
--
-- Unlike the JS variant, we do NOT collapse adjacent string literals via
-- `T.replace "' + '" ""`: that pattern incorrectly matches `' ' + ','`
-- (the space-then-comma sequence between a literal and `','.join(...)`),
-- producing `' ,'.join(...)` which uses ` ,` as the join separator and
-- swallows the leading space. The `intercalate " + "` output is correct
-- without further string fixups.
pySelfSyntaxText :: String -> TypeAndFields -> Expr -> Text
pySelfSyntaxText typeNamespace r expr =
  rewriteParams (paramAccessors r) (pySyntaxText typeNamespace r expr)

-- | Map field name to the Python access expression: `self['<name>']` for
-- required fields, `self.get('<name>')` for optional ones (since
-- TypedDict's `NotRequired` allows the key to be absent and `[...]` would
-- raise `KeyError`). Used by the rewriter so the same name is substituted
-- consistently in Optional `is not None` checks and in the value position.
paramAccessors :: TypeAndFields -> [(String, String)]
paramAccessors (_, fields) = map mk fields
  where
    mk (APIRecordField n t) = (n, accessor n t)
    accessor n = \case
      ATOptional _ -> "self.get('" ++ n ++ "')"
      _ -> "self['" ++ n ++ "']"

-- | Replace bare identifiers (matching a key in @accessors@) with the
-- corresponding accessor expression, skipping characters inside
-- single-quoted string literals and respecting identifier word boundaries.
rewriteParams :: [(String, String)] -> Text -> Text
rewriteParams accessors = T.pack . go False . T.unpack
  where
    go _ [] = []
    -- Toggle in/out of single-quoted string on every unescaped quote.
    go inStr ('\'' : rest) = '\'' : go (not inStr) rest
    go True (c : rest) = c : go True rest
    go False s@(c : rest)
      | isIdentStart c = case takeIdent s of
          (ident, after) -> case lookup ident accessors of
            Just expr -> expr ++ go False after
            Nothing -> ident ++ go False after
      | otherwise = c : go False rest
    isIdentStart c = isAlphaNum c || c == '_'
    takeIdent = span (\c -> isAlphaNum c || c == '_')