import os
from typing import Any, Dict, IO, Iterable, Type, Union, cast

from _internal_mitmproxy import exceptions
from _internal_mitmproxy import flow
from _internal_mitmproxy import flowfilter
from _internal_mitmproxy import http
from _internal_mitmproxy import tcp
from _internal_mitmproxy.io import compat
from _internal_mitmproxy.io import tnetstring

FLOW_TYPES: Dict[str, Type[flow.Flow]] = dict(
    http=http.HTTPFlow,
    tcp=tcp.TCPFlow,
)


class FlowWriter:
    def __init__(self, fo):
        self.fo = fo

    def add(self, flow):
        d = flow.get_state()
        tnetstring.dump(d, self.fo)


class FlowReader:
    def __init__(self, fo: IO[bytes]):
        self.fo: IO[bytes] = fo

    def stream(self) -> Iterable[flow.Flow]:
        """
            Yields Flow objects from the dump.
        """
        try:
            while True:
                # FIXME: This cast hides a lack of dynamic type checking
                loaded = cast(
                    Dict[Union[bytes, str], Any],
                    tnetstring.load(self.fo),
                )
                try:
                    mdata = compat.migrate_flow(loaded)
                except ValueError as e:
                    raise exceptions.FlowReadException(str(e))
                if mdata["type"] not in FLOW_TYPES:
                    raise exceptions.FlowReadException("Unknown flow type: {}".format(mdata["type"]))
                yield FLOW_TYPES[mdata["type"]].from_state(mdata)
        except (ValueError, TypeError, IndexError) as e:
            if str(e) == "not a tnetstring: empty file":
                return  # Error is due to EOF
            raise exceptions.FlowReadException("Invalid data format.")


class FilteredFlowWriter:
    def __init__(self, fo, flt):
        self.fo = fo
        self.flt = flt

    def add(self, f: flow.Flow):
        if self.flt and not flowfilter.match(self.flt, f):
            return
        d = f.get_state()
        tnetstring.dump(d, self.fo)


def read_flows_from_paths(paths):
    """
    Given a list of filepaths, read all flows and return a list of them.
    From a performance perspective, streaming would be advisable -
    however, if there's an error with one of the files, we want it to be raised immediately.

    Raises:
        FlowReadException, if any error occurs.
    """
    try:
        flows = []
        for path in paths:
            path = os.path.expanduser(path)
            with open(path, "rb") as f:
                flows.extend(FlowReader(f).stream())
    except OSError as e:
        raise exceptions.FlowReadException(e.strerror)
    return flows
