"""
Code generator: takes parsed TLObjects and generates Python source files
for tglib/tl/types/ and tglib/tl/functions/.
"""
import os
import struct
from collections import defaultdict
from pathlib import Path
from typing import List

from .parser import TLObject, TLArg

AUTO_GEN_NOTICE = '"""File auto-generated by tglib TL generator. Do not edit manually."""'

BASE_TYPES = {'string', 'bytes', 'int', 'long', 'int128', 'int256',
              'double', 'Bool', 'true', 'date'}


def _write_imports(f, depth: int, is_function: bool, tlobjects: List[TLObject]):
    """Write module-level imports."""
    base = '.' * depth
    f.write(f'{AUTO_GEN_NOTICE}\n')
    f.write(f'from {base}tl.tlobject import TLObject\n')
    if is_function:
        f.write(f'from {base}tl.tlobject import TLRequest\n')
    f.write('from typing import Optional, List, Union, TYPE_CHECKING\n')
    f.write('import os\n')
    f.write('import struct\n')
    f.write('from datetime import datetime\n')
    f.write('\n')


def _args_serialization(arg: TLArg, indent: str) -> List[str]:
    """Generate bytes serialization lines for one argument."""
    lines = []
    name = arg.name
    t = arg.type if not arg.flag_indicator else '#'

    def write(code):
        lines.append(f'{indent}{code}')

    if arg.flag_indicator:
        # Will be computed dynamically
        write(f'# flags computed below')
        return lines

    if arg.flag:
        # 'true' flag types are only reflected in the flags bitmask, never as fields
        if arg.type == 'true':
            return lines
        # Conditional field: only written if flag is set
        write(f'if self.{name} is not None:')
        sub = indent + '    '
        sub_lines = _serialize_single_type(name, arg.type, arg.is_vector,
                                           arg.use_vector_id,
                                           arg.skip_constructor_id, sub)
        if not sub_lines:
            lines.append(f'{sub}pass')
        else:
            lines.extend(sub_lines)
    else:
        sub_lines = _serialize_single_type(name, t, arg.is_vector,
                                           arg.use_vector_id,
                                           arg.skip_constructor_id, indent)
        lines.extend(sub_lines)

    return lines


def _serialize_single_type(name, t, is_vector, use_vector_id,
                            skip_constructor_id, indent) -> List[str]:
    lines = []

    def w(code):
        lines.append(f'{indent}{code}')

    if is_vector:
        if use_vector_id:
            w(f'buf.write(struct.pack(\'<i\', 0x1cb5c415))  # vector id')
        w(f'buf.write(struct.pack(\'<i\', len(self.{name})))')
        w(f'for item in self.{name}:')
        item_lines = _serialize_single_type('item', t, False, False,
                                            skip_constructor_id, indent + '    ')
        lines.extend(item_lines)
        return lines

    if t in ('int', 'date'):
        w(f'buf.write(struct.pack(\'<i\', self.{name}))')
    elif t == 'long':
        w(f'buf.write(struct.pack(\'<q\', self.{name}))')
    elif t == 'int128':
        w(f'buf.write(self.{name}.to_bytes(16, \'little\'))')
    elif t == 'int256':
        w(f'buf.write(self.{name}.to_bytes(32, \'little\'))')
    elif t == 'double':
        w(f'buf.write(struct.pack(\'<d\', self.{name}))')
    elif t in ('string', 'bytes'):
        w(f'buf.write(TLObject.serialize_bytes(self.{name}))')
    elif t == 'Bool':
        w(f'buf.write(struct.pack(\'<I\', 0x997275b5 if self.{name} else 0xbc799737))')
    elif t == 'true':
        pass  # true flags are only present/absent, no serialization needed
    else:
        if skip_constructor_id:
            w(f'buf.write(bytes(self.{name}))')
        else:
            w(f'buf.write(bytes(self.{name}))')

    return lines


def _write_to_bytes(obj: TLObject, f):
    """Write the _bytes() method."""
    indent = '        '
    f.write('    def _bytes(self) -> bytes:\n')
    f.write('        import io\n')
    f.write('        buf = io.BytesIO()\n')
    f.write(f'        buf.write(struct.pack(\'<I\', self.CONSTRUCTOR_ID))\n')

    # Handle flag computation
    has_flags = any(a.flag_indicator for a in obj.args)
    flag_fields = {}  # flag_name -> {bit_index: arg_name}
    for arg in obj.real_args:
        if arg.flag:
            flag_fields.setdefault(arg.flag, {})[arg.flag_index] = arg

    if has_flags:
        for flag_name, bits in flag_fields.items():
            f.write(f'        {flag_name} = 0\n')
            for bit, arg in bits.items():
                if arg.type == 'true':
                    f.write(f'        if self.{arg.name}:\n')
                    f.write(f'            {flag_name} |= (1 << {bit})\n')
                else:
                    f.write(f'        if self.{arg.name} is not None:\n')
                    f.write(f'            {flag_name} |= (1 << {bit})\n')

    for arg in obj.sorted_args() if hasattr(obj, 'sorted_args') else obj._sorted_args():
        if arg.generic_definition:
            continue
        if arg.flag_indicator:
            fname = arg.orig_name() if hasattr(arg, 'orig_name') else arg.name.rstrip('_').replace('is_self', 'self')
            f.write(f'        buf.write(struct.pack(\'<I\', {fname}))\n')
            continue

        lines = _args_serialization(arg, indent)
        for line in lines:
            f.write(line + '\n')

    f.write('        return buf.getvalue()\n\n')


def _write_from_reader(obj: TLObject, f):
    """Write the from_reader() classmethod."""
    f.write('    @classmethod\n')
    f.write('    def from_reader(cls, reader):\n')

    # Collect args
    real = [a for a in obj.args if not a.generic_definition]
    if not real:
        f.write('        return cls()\n\n')
        return

    has_flags = any(a.flag_indicator for a in real)
    flag_vars = set()

    f.write('        _args = {}\n')
    for arg in real:
        name = arg.name
        t = getattr(arg, 'type', None)

        if arg.flag_indicator:
            # Read flags int
            orig = arg.orig_name() if hasattr(arg, 'orig_name') else name
            f.write(f'        {orig} = reader.read_int(signed=False)\n')
            flag_vars.add(orig)
            continue

        if arg.flag:
            # Conditional read
            bit = arg.flag_index
            fname = arg.flag
            # 'true' flags only exist in the bitmask - no data to read
            if arg.type == 'true':
                f.write(f'        _args[\'{name}\'] = bool({fname} & (1 << {bit}))\n')
                continue
            f.write(f'        if {fname} & (1 << {bit}):\n')
            read_lines = _read_single_type(name, t, arg.is_vector,
                                           arg.use_vector_id,
                                           '            ')
            for line in read_lines:
                f.write(line + '\n')
            f.write(f'        else:\n')
            f.write(f'            _args[\'{name}\'] = None\n')
        else:
            read_lines = _read_single_type(name, t, arg.is_vector,
                                           arg.use_vector_id,
                                           '        ')
            for line in read_lines:
                f.write(line + '\n')

    f.write('        return cls(**_args)\n\n')


def _read_single_type(name, t, is_vector, use_vector_id, indent) -> List[str]:
    lines = []

    def w(code):
        lines.append(f'{indent}{code}')

    def assign(code):
        lines.append(f'{indent}_args[\'{name}\'] = {code}')

    if is_vector:
        if use_vector_id:
            w('reader.read_int(signed=False)  # skip vector id')
        w(f'_count_{name} = reader.read_int()')
        w(f'_list_{name} = []')
        w(f'for _ in range(_count_{name}):')
        item_lines = _read_single_item(f'_item_{name}', t, indent + '    ')
        lines.extend(item_lines)
        w(f'    _list_{name}.append(_item_{name})')
        assign(f'_list_{name}')
        return lines

    lines.extend(_read_single_item(f'_val_{name}', t, indent))
    assign(f'_val_{name}')
    return lines


def _read_single_item(var, t, indent) -> List[str]:
    lines = []

    def w(code):
        lines.append(f'{indent}{code}')

    if t in ('int', 'date'):
        w(f'{var} = reader.read_int()')
    elif t == 'long':
        w(f'{var} = reader.read_long()')
    elif t == 'int128':
        w(f'{var} = reader.read_large_int(128)')
    elif t == 'int256':
        w(f'{var} = reader.read_large_int(256)')
    elif t == 'double':
        w(f'{var} = reader.read_double()')
    elif t == 'string':
        w(f'{var} = reader.tgread_string()')
    elif t == 'bytes':
        w(f'{var} = reader.tgread_bytes()')
    elif t == 'Bool':
        w(f'{var} = reader.tgread_bool()')
    elif t == 'true':
        w(f'{var} = True')
    else:
        w(f'{var} = reader.tgread_object()')

    return lines


def _write_class(obj: TLObject, kind: str, f):
    """Write a complete class definition for a TLObject."""
    from zlib import crc32
    subclass_id = crc32(obj.result.encode('ascii')) & 0xFFFFFFFF

    f.write(f'\nclass {obj.class_name}({kind}):\n')
    f.write(f'    """TL type: {obj.fullname}"""\n')
    f.write(f'    CONSTRUCTOR_ID = {obj.id:#010x}\n')
    f.write(f'    SUBCLASS_OF_ID = {subclass_id:#010x}\n\n')

    # __init__
    if obj.real_args:
        # Build parameter list
        required = [a for a in obj.real_args if not a.flag and not a.can_be_inferred]
        optional = [a for a in obj.real_args if a.flag or a.can_be_inferred]

        params = []
        for a in required:
            params.append(f'{a.name}: {a.type_hint()}')
        for a in optional:
            params.append(f'{a.name}: {a.type_hint()} = None')

        params_str = ', '.join(params)
        f.write(f'    def __init__(self, {params_str}):\n')
        for a in obj.real_args:
            f.write(f'        self.{a.name} = {a.name}\n')
        f.write('\n')
    else:
        f.write('    def __init__(self):\n')
        f.write('        pass\n\n')

    # to_dict
    f.write('    def to_dict(self):\n')
    f.write(f'        return {{\n')
    f.write(f"            '_': '{obj.class_name}',\n")
    for a in obj.real_args:
        f.write(f"            '{a.name}': self.{a.name},\n")
    f.write('        }\n\n')

    # _bytes
    _write_to_bytes(obj, f)

    # from_reader
    _write_from_reader(obj, f)


def generate_tl_modules(
    tl_objects: List[TLObject],
    out_dir: str,
    depth: int = 2
):
    """
    Generate Python modules from TLObjects.

    Args:
        tl_objects: list from parse_tl()
        out_dir: root of tglib package (e.g. tglib/)
        depth: how many levels deep the generated files are from package root
    """
    types_dir = Path(out_dir) / 'tl' / 'types'
    funcs_dir = Path(out_dir) / 'tl' / 'functions'
    types_dir.mkdir(parents=True, exist_ok=True)
    funcs_dir.mkdir(parents=True, exist_ok=True)

    # Separate types and functions, group by namespace
    types_by_ns = defaultdict(list)
    funcs_by_ns = defaultdict(list)

    for obj in tl_objects:
        ns = obj.namespace or ''
        if obj.is_function:
            funcs_by_ns[ns].append(obj)
        else:
            types_by_ns[ns].append(obj)

    base = '.' * depth

    def write_module(ns_map, out_d, kind):
        all_classes = []
        ns_list = sorted(ns_map.keys())

        for ns in ns_list:
            objs = sorted(ns_map[ns], key=lambda x: x.name)
            filename = f'{ns}.py' if ns else '_root.py'
            filepath = out_d / filename

            with open(filepath, 'w', encoding='utf-8') as f:
                f.write(f'{AUTO_GEN_NOTICE}\n')
                f.write(f'from {base}tlobject import TLObject\n')
                if kind == 'TLRequest':
                    f.write(f'from {base}tlobject import TLRequest\n')
                f.write('from typing import Optional, List, Union\n')
                f.write('import os\nimport struct\nfrom datetime import datetime\n\n')

                for obj in objs:
                    _write_class(obj, kind, f)
                    all_classes.append((obj, ns))

            if ns:
                # Ensure namespace __init__.py exists
                ns_init = out_d / ns / '__init__.py'

        return ns_list, all_classes

    t_ns, t_classes = write_module(types_by_ns, types_dir, 'TLObject')
    f_ns, f_classes = write_module(funcs_by_ns, funcs_dir, 'TLRequest')

    # Generate types/__init__.py with namespace imports
    with open(types_dir / '__init__.py', 'w', encoding='utf-8') as f:
        f.write(f'{AUTO_GEN_NOTICE}\n')
        # Import root-namespace types from _root.py (avoids circular self-import)
        if '' in types_by_ns:
            objs = types_by_ns['']
            names = ', '.join(o.class_name for o in sorted(objs, key=lambda x: x.name))
            if names:
                f.write(f'from ._root import {names}  # noqa\n')
        for ns in sorted(n for n in types_by_ns if n):
            f.write(f'from . import {ns}  # noqa\n')

    with open(funcs_dir / '__init__.py', 'w', encoding='utf-8') as f:
        f.write(f'{AUTO_GEN_NOTICE}\n')
        # Import root-namespace functions from _root.py
        if '' in funcs_by_ns:
            objs = funcs_by_ns['']
            names = ', '.join(o.class_name for o in sorted(objs, key=lambda x: x.name))
            if names:
                f.write(f'from ._root import {names}  # noqa\n')
        for ns in sorted(n for n in funcs_by_ns if n):
            f.write(f'from . import {ns}  # noqa\n')

    # Generate alltlobjects.py
    all_objects_path = Path(out_dir) / 'tl' / 'alltlobjects.py'
    with open(all_objects_path, 'w', encoding='utf-8') as f:
        f.write(f'{AUTO_GEN_NOTICE}\n')
        f.write('# Maps constructor_id -> class for all TL types and functions\n\n')

        # imports for types
        for ns, objs in sorted(types_by_ns.items()):
            if ns:
                module = f'.types.{ns}'
            else:
                module = '.types._root'
            names = ', '.join(o.class_name for o in sorted(objs, key=lambda x: x.name))
            if names:
                f.write(f'from {module} import {names}\n')

        for ns, objs in sorted(funcs_by_ns.items()):
            if ns:
                module = f'.functions.{ns}'
            else:
                module = '.functions._root'
            names = ', '.join(o.class_name for o in sorted(objs, key=lambda x: x.name))
            if names:
                f.write(f'from {module} import {names}\n')

        f.write('\ntlobjects = {\n')
        for obj, ns in t_classes:
            f.write(f'    {obj.id:#010x}: {obj.class_name},\n')
        for obj, ns in f_classes:
            f.write(f'    {obj.id:#010x}: {obj.class_name},\n')
        f.write('}\n')

    print(f'✅ Generated {len(t_classes)} types and {len(f_classes)} functions')
    print(f'   Types dir:     {types_dir}')
    print(f'   Functions dir: {funcs_dir}')
    print(f'   All objects:   {all_objects_path}')
