import axios from 'axios'
import cheerio from 'cheerio'

const CODEFORCES_HOST = 'https://codeforces.com'

type ProblemDifficulty = 'A' | 'B' | 'C' | 'D' | 'E' | 'F'

export interface Example {
    input: string
    output: string
}

export class Problem {
    readonly name: string
    readonly problemSetId: number
    readonly difficulty: ProblemDifficulty

    private _examples: Example[] = []

    constructor(name: string, problemSetId: number, difficulty: ProblemDifficulty) {
        this.name = name;
        this.problemSetId = problemSetId;
        this.difficulty = difficulty;
    }

    get url(): string {
        return '/contest/' + this.problemSetId + '/problem/' + this.difficulty
    }

    get fullUrl(): string {
        return CODEFORCES_HOST + this.url
    }

    get examples() {
        return this._examples.slice()
    }

    async fetchExamples(): Promise<Example[]> {
        try {
            const problemSetResponse = await axios.get(this.url, { baseURL: CODEFORCES_HOST })
            if(problemSetResponse.status !== 200)
                throw Error('Error fetching input/output ' + this.problemSetId + '/' + this.difficulty)

            const problemHtml: string = problemSetResponse.data
            const $ = cheerio.load(problemHtml)

            const sampleTestElems = $('.problem-statement .sample-test > div')

            let sampleTests = []
            for(let i = 0; i < sampleTestElems.length; ++i) {
                
                // Add example only if output follows an input
                if(i + 1 < sampleTestElems.length && sampleTestElems[i + 1].attribs.class == 'output') {
                    sampleTests.push({
                        input: $('pre', sampleTestElems[i]).text(),
                        output: $('pre', sampleTestElems[i + 1]).text()
                    })
                    ++i
                }
            }

            this._examples = sampleTests
            return sampleTests
        }
        catch (e) {
            console.error(e)
            return []
        }
    }
}

export class ProblemSet {
    private _problems: Problem[] = []

    readonly id: number

    constructor(id: number) {
        this.id = id
    }

    get url(): string {
        return '/contest/' + this.id
    }

    get fullUrl(): string {
        return CODEFORCES_HOST + this.url
    }

    get problems(): Problem[] {
        return this._problems.slice()
    }

    async fetchAllExamples(): Promise<Problem[]> {
        for(const problem of this.problems) {
            await problem.fetchExamples()
        }
        return this.problems
    }

    async fetchProblems(): Promise<Problem[]> {
        try {
            const problemSetResponse = await axios.get(this.url, { baseURL: CODEFORCES_HOST })
            if(problemSetResponse.status !== 200)
                throw Error('Error fetching problem set ' + this.id)
            
            const problemSetHtml: string = problemSetResponse.data
            const $ = cheerio.load(problemSetHtml)

            const tableRows = $('.problems tbody tr')
            const problemRows = tableRows.slice(1, tableRows.length)

            this._problems = Array.from(problemRows).map((row) => {
                const links = Array.from($('a', row))
                const [ difficulty, name ] = links.map((link) => $(link).text().trim())

                return new Problem(name, this.id, difficulty as ProblemDifficulty)
            })

            return this.problems
        }
        catch (e) {
            console.error(e)
            return []
        }
    }
}