from typing import TYPE_CHECKING, cast


def optional_kv(key: str, value: str | None):
    if value:
        return f'{key}: {value}'
    return ''


def _prompt(input: str, class_name: str | None = None, user_hint: str | None = None):
    return rf"""
任务: 根据 Python 代码生成对应的 TypedDict 类型定义

{optional_kv('数据结构内容', class_name)}

{optional_kv('补充说明', user_hint)}

请根据以下输入的数据结构，生成对应的 python 接口代码，使用 TypedDict 定义类型。

输出要求：
- 参考模板样例生成模块代码
- 将生成的代码放入 <code/> 标签中
- 直接只输出 <code>...</code> ; 不包含其他任何标注、说明、解释等

### Input ###

输入的格式可能是 json, yml, key-value, 纯文本等各种包含结构化数据的格式，请你自行认真分析判断。

```
{input}
```

### TEMPLATE ###

```py
from typing import TypedDict, TYPE_CHECKING

if TYPE_CHECKING:
    # TODO add your type hints here
    ...

class [class_name](TypedDict):
    [fields]
```

### Example Output Format ###

<code>
from typing import TypedDict, TYPE_CHECKING, List
if TYPE_CHECKING:
    from numpy.typing import NDArray


class TimeseriesData(TypedDict):
    data: NDArray
    channels: List[str]
    srate: float
    t_second: float
</code>

"""


def gen_interface(
    input: str,
    class_name: str | None = None,
    user_hint: str | None = None,
    write_fpath: str | None = None,
    verbose=True,
):
    prompt = _prompt(input, class_name, user_hint)
    from zyplib.utils.openai import openai_complete

    result = openai_complete(
        prompt, return_type='text', stream=verbose, options={'temperature': 0.4}
    )

    if not isinstance(result, str):
        raise TypeError(f'openai_complete returned unexpected type: {type(result)}')

    result = result.strip()

    if result.startswith('<code>') and result.endswith('</code>'):
        result = result[6:-7]
    else:
        from zyplib.utils.print import print_debug

        print_debug(f'模板样例生成失败，请检查模板样例是否正确：\n{result}')
        return None

    if write_fpath:
        with open(write_fpath, 'w', encoding='utf-8') as f:
            f.write(result)
    return result
