import asyncio
import aiopg
import gc
import psycopg2
import psycopg2.extras
import socket
import random
import unittest
import time
import sys

from aiopg.connection import Connection, TIMEOUT
from aiopg.cursor import Cursor
from unittest import mock


PY_341 = sys.version_info >= (3, 4, 1)


class TestConnection(unittest.TestCase):

    def setUp(self):
        self.loop = asyncio.new_event_loop()
        asyncio.set_event_loop(None)

    def tearDown(self):
        self.loop.close()
        self.loop = None

    @asyncio.coroutine
    def connect(self, no_loop=False, **kwargs):
        loop = None if no_loop else self.loop
        conn = yield from aiopg.connect(database='aiopg',
                                        user='aiopg',
                                        password='passwd',
                                        host='127.0.0.1',
                                        loop=loop,
                                        **kwargs)
        conn2 = yield from aiopg.connect(database='aiopg',
                                         user='aiopg',
                                         password='passwd',
                                         host='127.0.0.1',
                                         loop=loop)
        cur = yield from conn2.cursor()
        yield from cur.execute("DROP TABLE IF EXISTS foo")
        yield from conn2.close()
        self.addCleanup(conn.close)
        return conn

    def test_connect(self):

        @asyncio.coroutine
        def go():
            conn = yield from self.connect()
            self.assertIsInstance(conn, Connection)
            self.assertFalse(conn._writing)
            self.assertIs(conn._conn, conn.raw)
            self.assertFalse(conn.echo)

        self.loop.run_until_complete(go())

    def test_simple_select(self):

        @asyncio.coroutine
        def go():
            conn = yield from self.connect()
            cur = yield from conn.cursor()
            self.assertIsInstance(cur, Cursor)
            yield from cur.execute('SELECT 1')
            ret = yield from cur.fetchone()
            self.assertEqual((1,), ret)

        self.loop.run_until_complete(go())

    def test_default_event_loop(self):
        asyncio.set_event_loop(self.loop)

        @asyncio.coroutine
        def go():
            conn = yield from self.connect(no_loop=True)
            cur = yield from conn.cursor()
            self.assertIsInstance(cur, Cursor)
            yield from cur.execute('SELECT 1')
            ret = yield from cur.fetchone()
            self.assertEqual((1,), ret)
            self.assertIs(conn._loop, self.loop)

        self.loop.run_until_complete(go())

    def test_close(self):

        @asyncio.coroutine
        def go():
            conn = yield from self.connect()
            yield from conn.close()
            self.assertTrue(conn.closed)

        self.loop.run_until_complete(go())

    def test_close_twice(self):

        @asyncio.coroutine
        def go():
            conn = yield from self.connect()
            yield from conn.close()
            yield from conn.close()
            self.assertTrue(conn.closed)

        self.loop.run_until_complete(go())

    def test_with_connection_factory(self):
        class Subclassed(aiopg.connection.Connection):
            pass

        @asyncio.coroutine
        def go():
            conn = yield from self.connect(
                aio_connection_factory=Subclassed,
            )
            self.assertIsInstance(conn, Subclassed)

        self.loop.run_until_complete(go())

    def test_with_cursor_factory(self):

        @asyncio.coroutine
        def go():
            conn = yield from self.connect()
            cur = yield from conn.cursor(
                cursor_factory=psycopg2.extras.DictCursor)
            yield from cur.execute('SELECT 1 AS a')
            ret = yield from cur.fetchone()
            self.assertEqual(1, ret['a'])

        self.loop.run_until_complete(go())

    def test_closed(self):

        @asyncio.coroutine
        def go():
            conn = yield from self.connect()
            self.assertFalse(conn.closed)
            yield from conn.close()
            self.assertTrue(conn.closed)

        self.loop.run_until_complete(go())

    def test_tpc(self):

        @asyncio.coroutine
        def go():
            conn = yield from self.connect()
            xid = yield from conn.xid(1, 'a', 'b')
            self.assertEqual((1, 'a', 'b'), tuple(xid))

            with self.assertRaises(psycopg2.ProgrammingError):
                yield from conn.tpc_begin(xid)
            with self.assertRaises(psycopg2.ProgrammingError):
                yield from conn.tpc_prepare()
            with self.assertRaises(psycopg2.ProgrammingError):
                yield from conn.tpc_commit(xid)
            with self.assertRaises(psycopg2.ProgrammingError):
                yield from conn.tpc_rollback(xid)
            with self.assertRaises(psycopg2.ProgrammingError):
                yield from conn.tpc_recover()

        self.loop.run_until_complete(go())

    def test_reset(self):

        @asyncio.coroutine
        def go():
            conn = yield from self.connect()

            with self.assertRaises(psycopg2.ProgrammingError):
                yield from conn.reset()

        self.loop.run_until_complete(go())

    def test_lobject(self):

        @asyncio.coroutine
        def go():
            conn = yield from self.connect()

            with self.assertRaises(psycopg2.ProgrammingError):
                yield from conn.lobject()

        self.loop.run_until_complete(go())

    def test_set_session(self):

        @asyncio.coroutine
        def go():
            conn = yield from self.connect()

            with self.assertRaises(psycopg2.ProgrammingError):
                yield from conn.set_session()

        self.loop.run_until_complete(go())

    def test_dsn(self):

        @asyncio.coroutine
        def go():
            conn = yield from self.connect()
            self.assertEqual(
                'dbname=aiopg user=aiopg password=xxxxxx host=127.0.0.1',
                conn.dsn)

        self.loop.run_until_complete(go())

    def test_get_backend_pid(self):

        @asyncio.coroutine
        def go():
            conn = yield from self.connect()

            ret = yield from conn.get_backend_pid()
            self.assertNotEqual(0, ret)

        self.loop.run_until_complete(go())

    def test_get_parameter_status(self):

        @asyncio.coroutine
        def go():
            conn = yield from self.connect()

            ret = yield from conn.get_parameter_status('is_superuser')
            self.assertEqual('off', ret)

        self.loop.run_until_complete(go())

    def test_cursor_factory(self):

        @asyncio.coroutine
        def go():
            conn = yield from self.connect(
                cursor_factory=psycopg2.extras.DictCursor)

            self.assertIs(psycopg2.extras.DictCursor, conn.cursor_factory)

        self.loop.run_until_complete(go())

    def test_notices(self):

        @asyncio.coroutine
        def go():
            conn = yield from self.connect()
            cur = yield from conn.cursor()
            yield from cur.execute("CREATE TABLE foo (id serial PRIMARY KEY);")

            if not conn.notices:
                raise unittest.SkipTest("Notices are disabled")

            self.assertEqual(
                ['NOTICE:  CREATE TABLE will create implicit sequence '
                 '"foo_id_seq" for serial column "foo.id"\n',
                 'NOTICE:  CREATE TABLE / PRIMARY KEY will create '
                 'implicit index "foo_pkey" for table "foo"\n'],
                conn.notices)

        self.loop.run_until_complete(go())

    def test_autocommit(self):

        @asyncio.coroutine
        def go():
            conn = yield from self.connect()

            self.assertTrue(conn.autocommit)