import asyncio
from contextlib import closing, redirect_stderr, redirect_stdout
from io import StringIO
import json
import os
import signal
from pathlib import Path
import sqlite3
import tempfile
import unittest
from unittest.mock import patch

from assembly_dataset.config import Config
from assembly_dataset.pipeline import generate, export, validate
from tests import test_assembly_pipeline as helpers


class InterruptedGenerationTests(unittest.TestCase):
    def test_completed_response_survives_cancel_and_resume(self):
        with tempfile.TemporaryDirectory() as tmp:
            config = helpers.PipelineTests().teacher_config(tmp)
            config.modes = ['branch']
            config.target_examples = 4
            class SlowAgent(helpers.TestAgent):
                calls = 0
                async def run(self, prompt, **kwargs):
                    self.calls += 1
                    if self.calls == 1:
                        await asyncio.Event().wait()
                    return await super().run(prompt, **kwargs)
            async def run():
                task = asyncio.create_task(generate(config, SlowAgent()))
                for _ in range(200):
                    await asyncio.sleep(0.01)
                    if (config.output_dir / 'episodes.sqlite3').exists():
                        with closing(sqlite3.connect(config.output_dir / 'episodes.sqlite3')) as db:
                            if db.execute('select count(*) from episodes').fetchone()[0]:
                                break
                task.cancel()
                with self.assertRaises(asyncio.CancelledError):
                    await task
            asyncio.run(run())
            with closing(sqlite3.connect(config.output_dir / 'episodes.sqlite3')) as db:
                self.assertEqual(db.execute('select count(*) from episodes').fetchone()[0], 1)
            self.assertEqual(export(config, allow_partial=True)['examples'], 1)
            self.assertEqual(validate(config)['examples_checked'], 1)
            agent = helpers.TestAgent()
            asyncio.run(generate(config, agent))
            self.assertEqual(export(config)['examples'], 4)
            self.assertEqual(validate(config)['episodes_checked'], 4)

    def test_api_failure_stops_without_symbolic_replacement(self):
        with tempfile.TemporaryDirectory() as tmp:
            config = helpers.PipelineTests().teacher_config(tmp)
            class FailedAgent:
                async def run(self, *args, **kwargs):
                    raise RuntimeError('quota exhausted secret')
            with self.assertRaisesRegex(RuntimeError, 'RuntimeError') as caught:
                asyncio.run(generate(config, FailedAgent()))
            self.assertNotIn('secret', str(caught.exception))
            with closing(sqlite3.connect(config.output_dir / 'episodes.sqlite3')) as db:
                self.assertEqual(db.execute('select count(*) from episodes').fetchone()[0], 0)

    def test_cli_exports_checkpoint_on_interrupt_or_error(self):
        from assembly_dataset.__main__ import main
        for failure, code in [(KeyboardInterrupt, 130), (RuntimeError, 1), (signal.SIGTERM, 130)]:
            with self.subTest(failure=failure), tempfile.TemporaryDirectory() as tmp:
                config = Config(output_dir=Path(tmp), target_examples=3, modes=["branch"])
                asyncio.run(generate(config))
                config.target_examples = 30
                async def interrupted(*args, **kwargs):
                    if failure == signal.SIGTERM:
                        os.kill(os.getpid(), signal.SIGTERM)
                        await asyncio.Event().wait()
                    else:
                        raise failure()
                with patch('sys.argv', ['assembly_dataset', '--no-progress']), patch('assembly_dataset.__main__.load_config', return_value=config), patch('assembly_dataset.__main__.generate', interrupted), redirect_stderr(StringIO()), redirect_stdout(StringIO()):
                    self.assertEqual(main(), code)
                manifest = json.loads((config.output_dir / 'sft/manifest.json').read_text())
                self.assertEqual(manifest['examples'], 3)
                self.assertTrue(manifest['partial'])
                self.assertEqual(validate(config)['examples_checked'], 3)

    def test_enough_rows_do_not_skip_interrupted_batch_holes(self):
        from assembly_dataset.pipeline import storage
        from assembly_dataset.world import make_episode, canonical
        with tempfile.TemporaryDirectory() as tmp:
            config = Config(output_dir=Path(tmp), target_examples=1, modes=['branch'])
            # Response for candidate 1 arrived; candidate 0 is still pending.
            episode = make_episode(1, config)
            with storage(config) as db, db:
                db.execute('INSERT INTO episodes VALUES (?,?,?,?)',
                           (2, episode['signature'], canonical(episode), 1))
            asyncio.run(generate(config))
            with storage(config) as db:
                indices = [json.loads(p)['candidate_index'] for (p,) in db.execute('select payload from episodes order by seq')]
            self.assertEqual(indices, [0, 1])

    def test_invalid_structured_response_keeps_validation_fallback(self):
        from pydantic_ai.exceptions import UnexpectedModelBehavior
        with tempfile.TemporaryDirectory() as tmp:
            config = helpers.PipelineTests().teacher_config(tmp)
            config.target_examples = 1
            class InvalidAgent:
                async def run(self, *args, **kwargs):
                    raise UnexpectedModelBehavior('Invalid structured output')
            asyncio.run(generate(config, InvalidAgent()))
            self.assertEqual(export(config)['teacher_failures'], 1)

    def test_teacher_only_rejects_symbolic_fallback(self):
        from assembly_dataset.pipeline import storage
        with tempfile.TemporaryDirectory() as tmp:
            config = helpers.PipelineTests().teacher_config(tmp)
            config.target_examples = 1
            with self.assertRaisesRegex(ValueError, 'teacher-only'):
                asyncio.run(generate(config, helpers.TestAgent(invalid=True), teacher_only=True))
            with storage(config) as db:
                self.assertEqual(db.execute('SELECT COUNT(*) FROM episodes').fetchone()[0], 0)

    def test_teacher_only_keeps_valid_teacher_response(self):
        with tempfile.TemporaryDirectory() as tmp:
            config = helpers.PipelineTests().teacher_config(tmp)
            config.target_examples = 1
            asyncio.run(generate(config, helpers.TestAgent(), teacher_only=True))
            self.assertEqual(export(config)['sources'], {'openai': 1})