import datetime
import os
import psycopg2
import redis
import socket

from flask import (
    abort,
    request,
    render_template,
    redirect,
    url_for,
    send_from_directory,
    send_file,
    g,
    jsonify,
    session,
)

from .models import (
    ModelForm,
    CASH_STRAT,
    CCY,
    BondDeal,
    CDSDeal,
    FUND,
    SwaptionDeal,
    FutureDeal,
    CashFlowDeal,
    CapFloorDeal,
    SpotDeal,
    Counterparties,
    Accounts,
    Termination,
)

from sqlalchemy.exc import IntegrityError
from sqlalchemy.sql import text
from sqlalchemy.sql.expression import func
from wtforms.fields import BooleanField

from .utils import bump_rev, simple_serialize
from PyPDF2 import PdfFileMerger
from io import BytesIO
from . import app
from . import db


def cp_choices(kind="bond"):
    if kind == "bond":
        return Counterparties.query.order_by("name").with_entities(
            Counterparties.code, Counterparties.name
        )
    elif kind in ["future", "spot"]:
        return []
    elif kind in ["cds", "swaption", "capfloor"]:
        return (
            Counterparties.query.order_by("name")
            .filter(Counterparties.name.ilike("%CDS%"))
            .with_entities(Counterparties.code, Counterparties.name)
        )


def account_codes():
    return Accounts.query.order_by("code").with_entities(Accounts.code, Accounts.name)


def fcm_accounts():
    return (
        Accounts.query.order_by("code")
        .filter(Accounts.name.ilike("%FCM%"))
        .with_entities(Accounts.code, Accounts.name)
    )


def get_queue():
    q = getattr(g, "queue", None)
    if q is None:
        hostname = socket.gethostname()
        if hostname == "ziggy":
            q = g.queue = redis.Redis(unix_socket_path="/run/redis/redis.sock")
        else:
            q = g.queue = redis.Redis(host="ziggy")
    return q


def get_db():
    db = getattr(g, "_database", None)
    if db is None:
        db = g._database = psycopg2.connect(
            database="serenitasdb", user="serenitas_user", host="debian"
        )
    return db


@app.teardown_appcontext
def close_connection(exception):
    db = getattr(g, "_database", None)
    if db is not None:
        db.close()


class CounterpartyForm(ModelForm):
    class Meta:
        model = Counterparties
        include_primary_keys = True


class BondForm(ModelForm):
    upload_globeop = BooleanField(label="Upload to globeop?")

    class Meta:
        model = BondDeal
        include_foreign_keys = True
        exclude = [
            "dealid",
            "lastupdate",  # we generate it with a trigger at the server level
            "principal_payment",
            "accrued_payment",
        ]


class CDSForm(ModelForm):
    upload_globeop = BooleanField(label="Upload to globeop?")

    class Meta:
        model = CDSDeal
        include_foreign_keys = True
        exclude = [
            "dealid",
            "lastupdate",
            "custodian",
            "cashaccount",
            "attach",
            "detach",
        ]


class SwaptionForm(ModelForm):
    upload_globeop = BooleanField(label="Upload to globeop?")

    class Meta:
        model = SwaptionDeal
        include_foreign_keys = True
        exclude = ["dealid", "lastupdate"]


class TerminationForm(ModelForm):
    upload_globeop = BooleanField(label="Upload to globeop?")

    class Meta:
        model = Termination
        include_foreign_keys = True


class FutureForm(ModelForm):
    upload_globeop = BooleanField(label="Upload to globeop?")

    class Meta:
        model = FutureDeal
        include_foreign_keys = True
        exclude = ["dealid", "lastupdate"]


class SpotForm(ModelForm):
    upload_globeop = BooleanField(label="Upload to globeop?")

    class Meta:
        model = SpotDeal
        include_foreign_keys = True
        exclude = ["dealid", "lastupdate"]


class CapFloorForm(ModelForm):
    upload_globeop = BooleanField(label="Upload to globeop?")

    class Meta:
        model = CapFloorDeal
        include_foreign_keys = True
        exclude = ["dealid", "lastupdate"]


def get_deal(kind):
    if kind == "cds":
        return CDSDeal
    elif kind == "bond":
        return BondDeal
    elif kind == "swaption":
        return SwaptionDeal
    elif kind == "future":
        return FutureDeal
    elif kind == "wire":
        return CashFlowDeal
    elif kind == "capfloor":
        return CapFloorDeal
    elif kind == "spot":
        return SpotDeal
    else:
        raise RuntimeError(f"Unknown Deal type: {kind}")


def _get_form(kind):
    if kind == "cds":
        return CDSForm
    elif kind == "bond":
        return BondForm
    elif kind == "swaption":
        return SwaptionForm
    elif kind == "future":
        return FutureForm
    elif kind == "capfloor":
        return CapFloorForm
    elif kind == "spot":
        return SpotForm
    else:
        raise RuntimeError("Unknown Deal type")


def get_form(trade, kind):
    Form = _get_form(kind)
    if trade.id:
        form = Form(obj=trade)
    else:
        form = Form()
        # add extra empty fields
        empty_choice = (None, "")
        for attr in ["folder", "buysell", "asset_class", "swaption_type"]:
            try:
                dropdown = getattr(form, attr)
                if dropdown.choices[0] != empty_choice:
                    dropdown.choices.insert(0, empty_choice)
            except AttributeError:
                continue

    if kind == "cds":
        form.account_code.choices = fcm_accounts()
        form.portfolio.choices = [
            c for c in form.portfolio.choices if c[0] not in ("IR", "IG", "HY")
        ]
        form.folder.choices = [
            c
            for c in form.folder.choices
            if (c[0] is None or (not c[0].startswith("SER_") or c[0].endswith("CURVE")))
        ]
    elif kind == "swaption":
        form.portfolio.choices = [("OPTIONS", "OPTIONS"), ("IR", "IR")]
    form.cp_code.choices = form.cp_code.choices + list(cp_choices(kind))
    return form


def get_trade(tradeid, kind):
    Deal = get_deal(kind)
    return Deal.query.get(tradeid) if tradeid else Deal()


def save_ticket(trade, old_ticket_name):
    if trade.ticket:
        if old_ticket_name:
            new_name = bump_rev(old_ticket_name)
        else:
            new_name = f"{trade.trade_date} {trade.id} {trade.description}.pdf"
        trade.ticket.save(os.path.join(app.config["TICKETS_FOLDER"], new_name))
        trade.ticket = new_name
    else:
        trade.ticket = old_ticket_name


def save_confirm(trade, old_confirm):
    d = {"C": "Cap", "F": "Floor"}
    if trade.trade_confirm:
        file_name = f"{trade.trade_date} {d[trade.cap_or_floor]}.pdf"
        trade.trade_confirm.save(os.path.join(app.config["CONFIRMS_FOLDER"], file_name))
        trade.trade_confirm = file_name
    else:
        trade.trade_confirm = old_confirm


def split_direction(g, direction):
    if direction == "outgoing":
        return [
            {
                "folder": cf.folder,
                "amount": -cf.amount,
                "code": cf.code,
                "currency": cf.currency,
                "action": cf.action,
            }
            for cf in g
            if cf.amount < 0
        ]
    elif direction == "incoming":
        return [
            {
                "folder": cf.folder,
                "amount": cf.amount,
                "code": cf.code,
                "currency": cf.currency,
                "action": cf.action,
            }
            for cf in g
            if cf.amount > 0
        ]
    else:
        raise ValueError("direction can be one of 'outgoing' or 'incoming'")


def gen_cashflow_deals(form, sql_session, wire_id=None):
    to_date = datetime.date.fromisoformat
    d = {
        "action": form.get("action"),
        "trade_date": form.get("trade_date", None, to_date),
    }

    for direction in ["incoming", "outgoing"]:
        count = 1
        while True:
            if f"{direction}-code-{count}" not in form:
                break
            else:
                d.update(
                    {
                        field: form.get(f"{direction}-{field}-{count}")
                        for field in ["folder", "code", "amount", "currency"]
                    }
                )
                count += 1
                if direction == "outgoing":
                    d["amount"] = -float(d["amount"])
                elif direction == "incoming":
                    d["amount"] = float(d["amount"])
                else:
                    raise ValueError("direction needs to be 'outgoing' or 'incoming'")
                if wire_id:
                    cf = CashFlowDeal.query.get(wire_id)
                    for k, v in d.items():
                        setattr(cf, k, v)
                    yield cf
                else:
                    cf = CashFlowDeal(**d)
                    sql_session.add(cf)
                    yield cf


@app.route("/wires/<int:wire_id>", methods=["GET", "POST"])
@app.route("/wires/", defaults={"wire_id": None}, methods=["GET", "POST"])
def wire_manage(wire_id):
    if request.method == "POST":
        wires = list(gen_cashflow_deals(request.form, db.session, wire_id))

        try:
            db.session.commit()
        except IntegrityError as e:
            app.logger.error(e)
            db.session.rollback()
            return render_template(
                "wire_entry.html",
                strategies=CASH_STRAT.enums,
                currencies=CCY.enums,
                accounts=account_codes(),
                outgoing_wires=split_direction(wires, "outgoing"),
                incoming_wires=split_direction(wires, "incoming"),
                action=request.form.get("action"),
                trade_date=request.form.get("trade_date"),
            )
        else:
            if request.form.get("upload_globeop") == "y":
                q = get_queue()
                for wire in wires:
                    q.rpush("wire_SERCGMAST", simple_serialize(wire))
            return redirect(url_for("list_trades", kind="wire"))

    wire = CashFlowDeal() if wire_id is None else CashFlowDeal.query.get(wire_id)
    return render_template(
        "wire_entry.html",
        strategies=CASH_STRAT.enums,
        currencies=CCY.enums,
        accounts=account_codes(),
        outgoing_wires=split_direction([wire], "outgoing") if wire_id else [],
        incoming_wires=split_direction([wire], "incoming") if wire_id else [],
        trade_date=wire.trade_date if wire_id else datetime.date.today(),
        action_url=url_for("wire_manage", wire_id=wire_id),
        action=wire.action if wire_id else None,
    )


@app.route("/trades/<kind>/<dealid>/terminate", methods=["GET", "POST"])
def terminate(dealid, kind):
    termination = Termination()
    form = TerminationForm(dealid=dealid)
    form.termination_cp.choices = form.termination_cp.choices + list(cp_choices(kind))
    table = kind if kind.endswith("s") else kind + "s"
    if form.validate_on_submit():
        form.populate_obj(termination)
        sql_session = form.get_session()
        rec = db.session.execute(
            "SELECT notional, coalesce(terminated_amount, 0.), currency, b.globeop_id, "
            "cp_code, fund "
            f"FROM {table} "
            "LEFT JOIN "
            "(SELECT dealid, sum(termination_amount) AS terminated_amount "
            " FROM terminations group by dealid) term USING (dealid) "
            "LEFT JOIN LATERAL ( "
            "  SELECT globeop_id FROM id_mapping "
            "     WHERE serenitas_id=id AND date <= :date "
            "  ORDER BY date DESC LIMIT 1"
            ") b ON true "
            "WHERE dealid = :dealid",
            {"dealid": dealid, "date": termination.termination_date},
        )
        notional, terminated_amount, currency, globeop_id, cp_code, fund = next(rec)
        is_assignment = True
        if not termination.partial_termination:
            termination.termination_amount = notional - terminated_amount
        if termination.termination_cp is None:
            termination.termination_cp = cp_code
            is_assignment = False
        sql_session.add(termination)
        try:
            sql_session.commit()
        except IntegrityError as e:
            app.logger.error(e)
            sql_session.rollback()
        else:
            buf = simple_serialize(
                termination,
                ccy=currency,
                globeopid=globeop_id,
                is_assignment=is_assignment,
                fund=fund,
            )
            q = get_queue()
            q.rpush(f"{kind}_termination", buf)
            return redirect(url_for("list_trades", kind=kind))
    else:
        return render_template(
            "termination.html",
            form=form,
            action_url=url_for("terminate", dealid=dealid, kind=kind),
        )


@app.route("/trades/<kind>/<int:tradeid>", methods=["GET", "POST"])
@app.route("/trades/<kind>/", defaults={"tradeid": None}, methods=["GET", "POST"])
@app.route(
    "/trades/", defaults={"tradeid": None, "kind": "bond"}, methods=["GET", "POST"]
)
def trade_manage(tradeid, kind):
    trade = get_trade(tradeid, kind)
    form = _get_form(kind)()

    if kind == "bond":
        old_ticket_name = trade.ticket
    if kind == "capfloor":
        old_confirm = trade.trade_confirm
    if kind == "cds":
        form.account_code.choices = fcm_accounts()
    form.cp_code.choices = form.cp_code.choices + list(cp_choices(kind))
    if form.validate_on_submit():
        form.populate_obj(trade)
        sql_session = form.get_session()
        if not tradeid:
            sql_session.add(trade)
        if kind == "bond":
            save_ticket(trade, old_ticket_name)
        if kind == "capfloor":
            save_confirm(trade, old_confirm)
        if kind == "cds":
            if trade.swap_type != "CD_INDEX":
                if trade.fund == "SERCGMAST":
                    trade.account_code = "BAC"
                elif trade.fund == "BRINKER":
                    trade.account_code = "BBH"
                elif trade.fund == "BOWDST":
                    trade.account_code = "BONY"
                else:
                    raise ValueError("Unknown fund")
            trade.cashaccount = trade.fcm_account.cash_account
            trade.custodian = trade.fcm_account.custodian
        try:
            sql_session.commit()
        except IntegrityError as e:
            app.logger.error(e)
            sql_session.rollback()
            return render_template(
                "trade_entry.html",
                form=form,
                action_url=url_for("trade_manage", tradeid=tradeid, kind=kind),
            )
        else:
            buf = simple_serialize(trade, upload=form.upload_globeop.data)
            q = get_queue()
            q.rpush(f"{kind}_{form.fund.data}", buf)
            return redirect(url_for("list_trades", kind=kind))
    else:
        if form.errors:
            app.logger.error(form.errors)
        form = get_form(trade, kind)
        return render_template(
            "trade_entry.html",
            form=form,
            action_url=url_for("trade_manage", tradeid=tradeid, kind=kind),
        )


@app.route("/", defaults={"kind": "bond", "fund": None})
@app.route("/<kind>/<fund>")
@app.route("/<kind>", defaults={"fund": None})
def list_trades(kind, fund):
    try:
        Deal = get_deal(kind)
    except RuntimeError as e:
        app.logger.error(e)
        abort(404)
    else:
        if kind not in ("cds", "swaption"):
            trade_list = Deal.query.order_by(Deal.trade_date.desc(), Deal.id.desc())
            if fund is not None:
                trade_list = trade_list.filter(
                    Deal.fund == func.cast(func.upper(fund), FUND)
                )
        else:
            if fund is not None:
                sql_str = text(
                    f"SELECT * FROM {kind}_trades WHERE fund=UPPER(:fund)::fund"
                )
                trade_list = db.session.execute(sql_str, {"fund": fund})
            else:
                trade_list = db.session.execute(f"SELECT * FROM {kind}_trades")
        return render_template(f"{kind}_blotter.html", trades=trade_list)


@app.route("/tickets/<int:tradeid>")
def download_ticket(tradeid):
    trade = BondDeal.query.get(tradeid)
    pdf = PdfFileMerger()
    pdf.append(os.path.join(app.config["TICKETS_FOLDER"], trade.ticket))
    pdf.append(os.path.join(app.config["CP_FOLDER"], trade.counterparty.instructions))
    fh = BytesIO()
    pdf.write(fh)
    pdf.close()
    fh.seek(0)
    return send_file(fh, mimetype="application/pdf")


@app.route("/confirms/<int:tradeid>")
def download_confirm(tradeid):
    trade = CapFloorDeal.query.get(tradeid)
    return send_file(
        os.path.join(app.config["CONFIRMS_FOLDER"], trade.trade_confirm),
        mimetype="application/pdf",
    )


@app.route("/counterparties/<path:instr>", methods=["GET"])
@app.route("/counterparties/", defaults={"instr": None}, methods=["GET"])
def list_counterparties(instr):
    if instr:
        return send_from_directory(
            filename=instr,
            directory=app.config["CP_FOLDER"],
            mimetype="application/pdf",
        )
    else:
        cp_list = Counterparties.query.order_by(Counterparties.name)
        return render_template("counterparties.html", counterparties=cp_list.all())


@app.route("/edit_cp/<cpcode>", methods=["GET", "POST"])
@app.route("/edit_cp/", defaults={"cpcode": None}, methods=["GET", "POST"])
def edit_counterparty(cpcode):
    if cpcode:
        cp = Counterparties.query.get(cpcode)
    else:
        cp = Counterparties()
    cp_form = CounterpartyForm()

    old_instructions = cp.instructions or None
    if cp_form.validate_on_submit():
        cp_form.populate_obj(cp)
        sql_session = cp_form.get_session()
        if not cpcode:
            sql_session.add(cp)
        instructions = cp_form.instructions
        if not instructions.data:
            cp.instructions = old_instructions
        else:
            cp.instructions = cp.name + ".pdf"
            instructions.data.save(
                os.path.join(app.config["CP_FOLDER"], cp.instructions)
            )
        sql_session.commit()
        return redirect(url_for("list_counterparties"))
    else:
        return render_template(
            "edit_cp.html", form=CounterpartyForm(obj=cp), code=cpcode
        )


@app.route("/_ajax", methods=["GET"])
def get_bbg_id():
    bbg_id = request.args.get("bbg_id")
    trade_date = request.args.get("trade_date", datetime.date.today())
    try:
        _, indextype, _, series, tenor = bbg_id.split()
    except ValueError:
        return "not a valid bloomberg description", 400
    indextype = indextype[:2]
    tenor = tenor[:-1] + "yr"
    series = int(series[1:])
    sql_str = ("SELECT redindexcode, maturity, coupon "
               "FROM index_desc "
               "WHERE index=%s and series=%s and tenor=%s "
               "      and lastdate >=%s ORDER BY version")
    db = get_db()
    with db.cursor() as c:
        c.execute(sql_str, (indextype, series, tenor, trade_date))
        redcode, maturity, coupon = c.fetchone()
    return jsonify(
        {
            "maturity": str(maturity),
            "redcode": redcode,
            "coupon": coupon,
        }
    )